registry.py 25.6 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
    "MiniMaxM1ForCausalLM": ("minimax_text_01", "MiniMaxText01ForCausalLM"),
40
41
42
43
    # 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
44
    "BambaForCausalLM": ("bamba", "BambaForCausalLM"),
45
    "BloomForCausalLM": ("bloom", "BloomForCausalLM"),
46
    "ChatGLMModel": ("chatglm", "ChatGLMForCausalLM"),
47
    "ChatGLMForConditionalGeneration": ("chatglm", "ChatGLMForCausalLM"),
48
    "CohereForCausalLM": ("commandr", "CohereForCausalLM"),
49
    "Cohere2ForCausalLM": ("commandr", "CohereForCausalLM"),
50
    "DbrxForCausalLM": ("dbrx", "DbrxForCausalLM"),
51
    "DeciLMForCausalLM": ("nemotron_nas", "DeciLMForCausalLM"),
52
53
    "DeepseekForCausalLM": ("deepseek", "DeepseekForCausalLM"),
    "DeepseekV2ForCausalLM": ("deepseek_v2", "DeepseekV2ForCausalLM"),
54
    "DeepseekV3ForCausalLM": ("deepseek_v2", "DeepseekV3ForCausalLM"),
55
56
    "ExaoneForCausalLM": ("exaone", "ExaoneForCausalLM"),
    "FalconForCausalLM": ("falcon", "FalconForCausalLM"),
57
    "FM9GForCausalLM": ("fm9g", "FM9GForCausalLM"),
58
    "Fairseq2LlamaForCausalLM": ("fairseq2_llama", "Fairseq2LlamaForCausalLM"),
59
60
    "GemmaForCausalLM": ("gemma", "GemmaForCausalLM"),
    "Gemma2ForCausalLM": ("gemma2", "Gemma2ForCausalLM"),
61
    "Gemma3ForCausalLM": ("gemma3", "Gemma3ForCausalLM"),
62
    "GlmForCausalLM": ("glm", "GlmForCausalLM"),
Yuxuan Zhang's avatar
Yuxuan Zhang committed
63
    "Glm4ForCausalLM": ("glm4", "Glm4ForCausalLM"),
64
65
66
67
68
69
    "GPT2LMHeadModel": ("gpt2", "GPT2LMHeadModel"),
    "GPTBigCodeForCausalLM": ("gpt_bigcode", "GPTBigCodeForCausalLM"),
    "GPTJForCausalLM": ("gpt_j", "GPTJForCausalLM"),
    "GPTNeoXForCausalLM": ("gpt_neox", "GPTNeoXForCausalLM"),
    "GraniteForCausalLM": ("granite", "GraniteForCausalLM"),
    "GraniteMoeForCausalLM": ("granitemoe", "GraniteMoeForCausalLM"),
70
    "GraniteMoeHybridForCausalLM": ("granitemoehybrid", "GraniteMoeHybridForCausalLM"),   # noqa: E501
71
    "GraniteMoeSharedForCausalLM": ("granitemoeshared", "GraniteMoeSharedForCausalLM"),   # noqa: E501
72
    "GritLM": ("gritlm", "GritLM"),
Michael Goin's avatar
Michael Goin committed
73
    "Grok1ModelForCausalLM": ("grok1", "Grok1ForCausalLM"),
74
75
    "InternLMForCausalLM": ("llama", "LlamaForCausalLM"),
    "InternLM2ForCausalLM": ("internlm2", "InternLM2ForCausalLM"),
76
    "InternLM2VEForCausalLM": ("internlm2_ve", "InternLM2VEForCausalLM"),
77
    "InternLM3ForCausalLM": ("llama", "LlamaForCausalLM"),
78
79
80
    "JAISLMHeadModel": ("jais", "JAISLMHeadModel"),
    "JambaForCausalLM": ("jamba", "JambaForCausalLM"),
    "LlamaForCausalLM": ("llama", "LlamaForCausalLM"),
81
    "Llama4ForCausalLM": ("llama4", "Llama4ForCausalLM"),
82
83
    # For decapoda-research/llama-*
    "LLaMAForCausalLM": ("llama", "LlamaForCausalLM"),
84
    "MambaForCausalLM": ("mamba", "MambaForCausalLM"),
85
    "FalconMambaForCausalLM": ("mamba", "MambaForCausalLM"),
Dhia Eddine Rhaiem's avatar
Dhia Eddine Rhaiem committed
86
    "FalconH1ForCausalLM":("falcon_h1", "FalconH1ForCausalLM"),
87
    "Mamba2ForCausalLM": ("mamba2", "Mamba2ForCausalLM"),
88
89
    "MiniCPMForCausalLM": ("minicpm", "MiniCPMForCausalLM"),
    "MiniCPM3ForCausalLM": ("minicpm3", "MiniCPM3ForCausalLM"),
90
91
92
93
94
95
    "MistralForCausalLM": ("llama", "LlamaForCausalLM"),
    "MixtralForCausalLM": ("mixtral", "MixtralForCausalLM"),
    "QuantMixtralForCausalLM": ("mixtral_quant", "MixtralForCausalLM"),
    # transformers's mpt class has lower case
    "MptForCausalLM": ("mpt", "MPTForCausalLM"),
    "MPTForCausalLM": ("mpt", "MPTForCausalLM"),
96
    "MiMoForCausalLM": ("mimo", "MiMoForCausalLM"),
97
    "NemotronForCausalLM": ("nemotron", "NemotronForCausalLM"),
Luis Vega's avatar
Luis Vega committed
98
    "NemotronHForCausalLM": ("nemotron_h", "NemotronHForCausalLM"),
99
    "OlmoForCausalLM": ("olmo", "OlmoForCausalLM"),
100
    "Olmo2ForCausalLM": ("olmo2", "Olmo2ForCausalLM"),
101
102
103
104
105
106
107
108
    "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
109
    "Plamo2ForCausalLM": ("plamo2", "Plamo2ForCausalLM"),
110
    "QWenLMHeadModel": ("qwen", "QWenLMHeadModel"),
111
112
    "Qwen2ForCausalLM": ("qwen2", "Qwen2ForCausalLM"),
    "Qwen2MoeForCausalLM": ("qwen2_moe", "Qwen2MoeForCausalLM"),
113
114
    "Qwen3ForCausalLM": ("qwen3", "Qwen3ForCausalLM"),
    "Qwen3MoeForCausalLM": ("qwen3_moe", "Qwen3MoeForCausalLM"),
115
116
117
118
119
    "RWForCausalLM": ("falcon", "FalconForCausalLM"),
    "StableLMEpochForCausalLM": ("stablelm", "StablelmForCausalLM"),
    "StableLmForCausalLM": ("stablelm", "StablelmForCausalLM"),
    "Starcoder2ForCausalLM": ("starcoder2", "Starcoder2ForCausalLM"),
    "SolarForCausalLM": ("solar", "SolarForCausalLM"),
120
    "TeleChat2ForCausalLM": ("telechat2", "TeleChat2ForCausalLM"),
121
    "TeleFLMForCausalLM": ("teleflm", "TeleFLMForCausalLM"),
122
    "XverseForCausalLM": ("llama", "LlamaForCausalLM"),
123
    "Zamba2ForCausalLM": ("zamba2", "Zamba2ForCausalLM"),
124
125
    "Ernie4_5_ForCausalLM": ("ernie45", "Ernie4_5_ForCausalLM"),
    "Ernie4_5_MoeForCausalLM": ("ernie45_moe", "Ernie4_5_MoeForCausalLM"),
126
127
128
    # [Encoder-decoder]
    "BartModel": ("bart", "BartForConditionalGeneration"),
    "BartForConditionalGeneration": ("bart", "BartForConditionalGeneration"),
129
130
131
}

_EMBEDDING_MODELS = {
132
    # [Text-only]
133
    "BertModel": ("bert", "BertEmbeddingModel"),
134
    "DeciLMForCausalLM": ("nemotron_nas", "DeciLMForCausalLM"),
135
    "Gemma2Model": ("gemma2", "Gemma2ForCausalLM"),
136
    "GlmForCausalLM": ("glm", "GlmForCausalLM"),
137
    "GritLM": ("gritlm", "GritLM"),
138
139
    "GteModel": ("bert_with_rope", "SnowflakeGteNewModel"),
    "GteNewModel": ("bert_with_rope", "GteNewModel"),
140
    "InternLM2ForRewardModel": ("internlm2", "InternLM2ForRewardModel"),
141
    "JambaForSequenceClassification": ("jamba", "JambaForSequenceClassification"),  # noqa: E501
142
    "LlamaModel": ("llama", "LlamaForCausalLM"),
143
144
145
146
147
    **{
        # 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"
    },
148
    "MistralModel": ("llama", "LlamaForCausalLM"),
149
    "ModernBertModel": ("modernbert", "ModernBertModel"),
150
    "NomicBertModel": ("bert_with_rope", "NomicBertModel"),
151
    "Phi3ForCausalLM": ("phi3", "Phi3ForCausalLM"),
152
    "Qwen2Model": ("qwen2", "Qwen2ForCausalLM"),
153
    "Qwen2ForCausalLM": ("qwen2", "Qwen2ForCausalLM"),
154
    "Qwen2ForRewardModel": ("qwen2_rm", "Qwen2ForRewardModel"),
155
    "Qwen2ForProcessRewardModel": ("qwen2_rm", "Qwen2ForProcessRewardModel"),
156
157
    "RobertaForMaskedLM": ("roberta", "RobertaEmbeddingModel"),
    "RobertaModel": ("roberta", "RobertaEmbeddingModel"),
158
    "TeleChat2ForCausalLM": ("telechat2", "TeleChat2ForCausalLM"),
159
    "XLMRobertaModel": ("roberta", "RobertaEmbeddingModel"),
160
    # [Multimodal]
Cyrus Leung's avatar
Cyrus Leung committed
161
    "LlavaNextForConditionalGeneration": ("llava_next", "LlavaNextForConditionalGeneration"),  # noqa: E501
162
    "Phi3VForCausalLM": ("phi3v", "Phi3VForCausalLM"),
163
    "Qwen2VLForConditionalGeneration": ("qwen2_vl", "Qwen2VLForConditionalGeneration"),  # noqa: E501
164
165
    # [Auto-converted (see adapters.py)]
    "Qwen2ForSequenceClassification": ("qwen2", "Qwen2ForCausalLM"),
166
167
168
169
    # 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"),
170
171
}

172
173
174
175
176
177
_CROSS_ENCODER_MODELS = {
    "BertForSequenceClassification": ("bert", "BertForSequenceClassification"),
    "RobertaForSequenceClassification": ("roberta",
                                         "RobertaForSequenceClassification"),
    "XLMRobertaForSequenceClassification": ("roberta",
                                            "RobertaForSequenceClassification"),
xsank's avatar
xsank committed
178
179
    "ModernBertForSequenceClassification": ("modernbert",
                                            "ModernBertForSequenceClassification"),
180
    "Qwen3ForSequenceClassification": ("qwen3", "Qwen3ForSequenceClassification"), # noqa: E501
181
182
}

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

_SPECULATIVE_DECODING_MODELS = {
232
    "MiMoMTPModel": ("mimo_mtp", "MiMoMTP"),
233
    "EAGLEModel": ("eagle", "EAGLE"),
234
    "EagleLlamaForCausalLM": ("llama_eagle", "EagleLlamaForCausalLM"),
235
    "EagleMiniCPMForCausalLM": ("minicpm_eagle", "EagleMiniCPMForCausalLM"),
236
    "Eagle3LlamaForCausalLM": ("llama_eagle3", "Eagle3LlamaForCausalLM"),
237
    "DeepSeekMTPModel": ("deepseek_mtp", "DeepSeekMTP"),
238
239
    "MedusaModel": ("medusa", "Medusa"),
    "MLPSpeculatorPreTrainedModel": ("mlp_speculator", "MLPSpeculator"),
240
}
241

242
_TRANSFORMERS_MODELS = {
243
    "TransformersForCausalLM": ("transformers", "TransformersForCausalLM"),
244
}
245
# yapf: enable
246

247
_VLLM_MODELS = {
248
    **_TEXT_GENERATION_MODELS,
249
    **_EMBEDDING_MODELS,
250
    **_CROSS_ENCODER_MODELS,
251
    **_MULTIMODAL_MODELS,
252
    **_SPECULATIVE_DECODING_MODELS,
253
    **_TRANSFORMERS_MODELS,
254
255
}

256
257
258
259
260
261
262
263
# 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"
]

264

265
266
@dataclass(frozen=True)
class _ModelInfo:
267
    architecture: str
268
    is_text_generation_model: bool
269
    is_pooling_model: bool
270
    supports_cross_encoding: bool
271
272
    supports_multimodal: bool
    supports_pp: bool
273
274
    has_inner_state: bool
    is_attention_free: bool
275
    is_hybrid: bool
276
    has_noops: bool
277
    supports_transcription: bool
278
    supports_v0_only: bool
279
280

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


298
class _BaseRegisteredModel(ABC):
299

300
301
302
    @abstractmethod
    def inspect_model_cls(self) -> _ModelInfo:
        raise NotImplementedError
303

304
    @abstractmethod
305
    def load_model_cls(self) -> type[nn.Module]:
306
        raise NotImplementedError
307
308


309
310
311
312
313
314
315
@dataclass(frozen=True)
class _RegisteredModel(_BaseRegisteredModel):
    """
    Represents a model that has already been imported in the main process.
    """

    interfaces: _ModelInfo
316
    model_cls: type[nn.Module]
317
318

    @staticmethod
319
    def from_model_cls(model_cls: type[nn.Module]):
320
321
322
323
324
325
326
327
        return _RegisteredModel(
            interfaces=_ModelInfo.from_model_cls(model_cls),
            model_cls=model_cls,
        )

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

328
    def load_model_cls(self) -> type[nn.Module]:
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
        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()))

345
    def load_model_cls(self) -> type[nn.Module]:
346
347
348
349
350
351
352
353
        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,
354
) -> Optional[type[nn.Module]]:
355
    from vllm.platforms import current_platform
356
    current_platform.verify_model_arch(model_arch)
357
358
359
360
361
362
    try:
        return model.load_model_cls()
    except Exception:
        logger.exception("Error in loading model architecture '%s'",
                         model_arch)
        return None
363
364


365
366
367
368
369
370
371
372
373
374
375
@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
376
377


378
379
380
@dataclass
class _ModelRegistry:
    # Keyed by model_arch
381
    models: dict[str, _BaseRegisteredModel] = field(default_factory=dict)
382

383
    def get_supported_archs(self) -> Set[str]:
384
        return self.models.keys()
385

386
387
388
    def register_model(
        self,
        model_arch: str,
389
        model_cls: Union[type[nn.Module], str],
390
    ) -> None:
391
392
393
        """
        Register an external model to be used in vLLM.

394
        `model_cls` can be either:
395

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

406
        if model_arch in self.models:
407
408
409
            logger.warning(
                "Model architecture %s is already registered, and will be "
                "overwritten by the new model class %s.", model_arch,
410
411
412
413
414
415
416
                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)
417

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

426
        self.models[model_arch] = model
427

428
    def _raise_for_unsupported(self, architectures: list[str]):
429
        all_supported_archs = self.get_supported_archs()
430

431
432
433
434
435
        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.")

436
437
438
        raise ValueError(
            f"Model architectures {architectures} are not supported for now. "
            f"Supported architectures: {all_supported_archs}")
439

440
    def _try_load_model_cls(self,
441
                            model_arch: str) -> Optional[type[nn.Module]]:
442
443
        if model_arch not in self.models:
            return None
444

445
        return _try_load_model_cls(model_arch, self.models[model_arch])
446

447
448
449
    def _try_inspect_model_cls(self, model_arch: str) -> Optional[_ModelInfo]:
        if model_arch not in self.models:
            return None
450

451
        return _try_inspect_model_cls(model_arch, self.models[model_arch])
452

453
454
    def _normalize_archs(
        self,
455
456
        architectures: Union[str, list[str]],
    ) -> list[str]:
457
458
459
460
461
        if isinstance(architectures, str):
            architectures = [architectures]
        if not architectures:
            logger.warning("No model architectures are specified")

462
463
464
465
        # filter out support architectures
        normalized_arch = list(
            filter(lambda model: model in self.models, architectures))

466
        # make sure Transformers backend is put at the last as a fallback
467
        if len(normalized_arch) != len(architectures):
468
            normalized_arch.append("TransformersForCausalLM")
469
        return normalized_arch
470

471
472
    def inspect_model_cls(
        self,
473
474
        architectures: Union[str, list[str]],
    ) -> tuple[_ModelInfo, str]:
475
        architectures = self._normalize_archs(architectures)
476

477
478
479
        for arch in architectures:
            model_info = self._try_inspect_model_cls(arch)
            if model_info is not None:
480
                return (model_info, arch)
481

482
        return self._raise_for_unsupported(architectures)
483

484
485
    def resolve_model_cls(
        self,
486
487
        architectures: Union[str, list[str]],
    ) -> tuple[type[nn.Module], str]:
488
        architectures = self._normalize_archs(architectures)
489

490
491
492
493
        for arch in architectures:
            model_cls = self._try_load_model_cls(arch)
            if model_cls is not None:
                return (model_cls, arch)
494

495
        return self._raise_for_unsupported(architectures)
496

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

504
    def is_pooling_model(
505
        self,
506
        architectures: Union[str, list[str]],
507
    ) -> bool:
508
        model_cls, _ = self.inspect_model_cls(architectures)
509
        return model_cls.is_pooling_model
510

511
512
    def is_cross_encoder_model(
        self,
513
        architectures: Union[str, list[str]],
514
    ) -> bool:
515
516
        model_cls, _ = self.inspect_model_cls(architectures)
        return model_cls.supports_cross_encoding
517

518
519
    def is_multimodal_model(
        self,
520
        architectures: Union[str, list[str]],
521
    ) -> bool:
522
523
        model_cls, _ = self.inspect_model_cls(architectures)
        return model_cls.supports_multimodal
524
525
526

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

532
533
    def model_has_inner_state(
        self,
534
        architectures: Union[str, list[str]],
535
536
537
    ) -> bool:
        model_cls, _ = self.inspect_model_cls(architectures)
        return model_cls.has_inner_state
538

539
540
    def is_attention_free_model(
        self,
541
        architectures: Union[str, list[str]],
542
543
544
    ) -> bool:
        model_cls, _ = self.inspect_model_cls(architectures)
        return model_cls.is_attention_free
545

546
547
    def is_hybrid_model(
        self,
548
        architectures: Union[str, list[str]],
549
550
551
552
    ) -> bool:
        model_cls, _ = self.inspect_model_cls(architectures)
        return model_cls.is_hybrid

553
554
    def is_noops_model(
        self,
555
        architectures: Union[str, list[str]],
556
557
558
559
    ) -> bool:
        model_cls, _ = self.inspect_model_cls(architectures)
        return model_cls.has_noops

560
561
    def is_transcription_model(
        self,
562
        architectures: Union[str, list[str]],
563
564
565
566
    ) -> bool:
        model_cls, _ = self.inspect_model_cls(architectures)
        return model_cls.supports_transcription

567
568
    def is_v1_compatible(
        self,
569
        architectures: Union[str, list[str]],
570
571
572
573
    ) -> bool:
        model_cls, _ = self.inspect_model_cls(architectures)
        return not model_cls.supports_v0_only

574
575

ModelRegistry = _ModelRegistry({
576
577
    model_arch:
    _LazyRegisteredModel(
578
579
580
581
582
583
584
585
586
587
        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:
588
589
590
591
592
    # 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")

593
        # `cloudpickle` allows pickling lambda functions directly
594
        input_bytes = cloudpickle.dumps((fn, output_filepath))
595
596
597

        # cannot use `sys.executable __file__` here because the script
        # contains relative imports
598
599
600
        returned = subprocess.run(_SUBPROCESS_COMMAND,
                                  input=input_bytes,
                                  capture_output=True)
601
602
603
604
605
606
607
608
609

        # 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

610
        with open(output_filepath, "rb") as f:
611
612
613
614
615
616
617
618
619
620
621
            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()
622
623
624

    with open(output_file, "wb") as f:
        f.write(pickle.dumps(result))
625
626
627


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