registry.py 38 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
from dataclasses import dataclass, field
16
from functools import lru_cache
17
from typing import Callable, Optional, TypeVar, Union
18
19

import torch.nn as nn
20
import transformers
21

22
23
from vllm.config import (ModelConfig, ModelImpl, iter_architecture_defaults,
                         try_match_architecture_defaults)
24
from vllm.logger import init_logger
25
26
from vllm.transformers_utils.dynamic_module import (
    try_get_class_from_dynamic_module)
27

28
29
from .interfaces import (has_inner_state, has_noops, is_attention_free,
                         is_hybrid, supports_cross_encoding,
30
31
                         supports_multimodal,
                         supports_multimodal_encoder_tp_data,
32
                         supports_multimodal_raw_input_only, supports_pp,
33
                         supports_transcription, supports_v0_only)
34
35
from .interfaces_base import (get_default_pooling_type, is_pooling_model,
                              is_text_generation_model)
36
37
38

logger = init_logger(__name__)

39
# yapf: disable
40
41
_TEXT_GENERATION_MODELS = {
    # [Decoder-only]
42
43
    "AquilaModel": ("llama", "LlamaForCausalLM"),
    "AquilaForCausalLM": ("llama", "LlamaForCausalLM"),  # AquilaChat2
Raghav Ravishankar's avatar
Raghav Ravishankar committed
44
    "ArceeForCausalLM": ("arcee", "ArceeForCausalLM"),
45
    "ArcticForCausalLM": ("arctic", "ArcticForCausalLM"),
46
    "MiniMaxForCausalLM": ("minimax_text_01", "MiniMaxText01ForCausalLM"),
47
    "MiniMaxText01ForCausalLM": ("minimax_text_01", "MiniMaxText01ForCausalLM"),
48
    "MiniMaxM1ForCausalLM": ("minimax_text_01", "MiniMaxText01ForCausalLM"),
49
50
51
52
    # baichuan-7b, upper case 'C' in the class name
    "BaiChuanForCausalLM": ("baichuan", "BaiChuanForCausalLM"),
    # baichuan-13b, lower case 'c' in the class name
    "BaichuanForCausalLM": ("baichuan", "BaichuanForCausalLM"),
53
    "BailingMoeForCausalLM": ("bailing_moe", "BailingMoeForCausalLM"),
Yu Chin Fabian Lim's avatar
Yu Chin Fabian Lim committed
54
    "BambaForCausalLM": ("bamba", "BambaForCausalLM"),
55
    "BloomForCausalLM": ("bloom", "BloomForCausalLM"),
56
    "ChatGLMModel": ("chatglm", "ChatGLMForCausalLM"),
57
    "ChatGLMForConditionalGeneration": ("chatglm", "ChatGLMForCausalLM"),
58
    "CohereForCausalLM": ("commandr", "CohereForCausalLM"),
59
    "Cohere2ForCausalLM": ("commandr", "CohereForCausalLM"),
60
    "DbrxForCausalLM": ("dbrx", "DbrxForCausalLM"),
61
    "DeciLMForCausalLM": ("nemotron_nas", "DeciLMForCausalLM"),
62
63
    "DeepseekForCausalLM": ("deepseek", "DeepseekForCausalLM"),
    "DeepseekV2ForCausalLM": ("deepseek_v2", "DeepseekV2ForCausalLM"),
64
    "DeepseekV3ForCausalLM": ("deepseek_v2", "DeepseekV3ForCausalLM"),
65
    "Dots1ForCausalLM": ("dots1", "Dots1ForCausalLM"),
66
    "Ernie4_5ForCausalLM": ("ernie45", "Ernie4_5ForCausalLM"),
67
    "Ernie4_5_MoeForCausalLM": ("ernie45_moe", "Ernie4_5_MoeForCausalLM"),
68
    "ExaoneForCausalLM": ("exaone", "ExaoneForCausalLM"),
69
    "Exaone4ForCausalLM": ("exaone4", "Exaone4ForCausalLM"),
70
    "FalconForCausalLM": ("falcon", "FalconForCausalLM"),
71
    "Fairseq2LlamaForCausalLM": ("fairseq2_llama", "Fairseq2LlamaForCausalLM"),
72
73
    "GemmaForCausalLM": ("gemma", "GemmaForCausalLM"),
    "Gemma2ForCausalLM": ("gemma2", "Gemma2ForCausalLM"),
74
    "Gemma3ForCausalLM": ("gemma3", "Gemma3ForCausalLM"),
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
75
    "Gemma3nForCausalLM": ("gemma3n", "Gemma3nForCausalLM"),
76
    "GlmForCausalLM": ("glm", "GlmForCausalLM"),
Yuxuan Zhang's avatar
Yuxuan Zhang committed
77
    "Glm4ForCausalLM": ("glm4", "Glm4ForCausalLM"),
Yuxuan Zhang's avatar
Yuxuan Zhang committed
78
    "Glm4MoeForCausalLM": ("glm4_moe", "Glm4MoeForCausalLM"),
79
    "GptOssForCausalLM": ("gpt_oss", "GptOssForCausalLM"),
80
81
82
83
84
85
    "GPT2LMHeadModel": ("gpt2", "GPT2LMHeadModel"),
    "GPTBigCodeForCausalLM": ("gpt_bigcode", "GPTBigCodeForCausalLM"),
    "GPTJForCausalLM": ("gpt_j", "GPTJForCausalLM"),
    "GPTNeoXForCausalLM": ("gpt_neox", "GPTNeoXForCausalLM"),
    "GraniteForCausalLM": ("granite", "GraniteForCausalLM"),
    "GraniteMoeForCausalLM": ("granitemoe", "GraniteMoeForCausalLM"),
86
    "GraniteMoeHybridForCausalLM": ("granitemoehybrid", "GraniteMoeHybridForCausalLM"),   # noqa: E501
87
    "GraniteMoeSharedForCausalLM": ("granitemoeshared", "GraniteMoeSharedForCausalLM"),   # noqa: E501
88
    "GritLM": ("gritlm", "GritLM"),
Michael Goin's avatar
Michael Goin committed
89
    "Grok1ModelForCausalLM": ("grok1", "Grok1ForCausalLM"),
90
91
    "HunYuanMoEV1ForCausalLM": ("hunyuan_v1", "HunYuanMoEV1ForCausalLM"),
    "HunYuanDenseV1ForCausalLM": ("hunyuan_v1", "HunYuanDenseV1ForCausalLM"),
92
    "HCXVisionForCausalLM": ("hyperclovax_vision", "HCXVisionForCausalLM"),
93
94
    "InternLMForCausalLM": ("llama", "LlamaForCausalLM"),
    "InternLM2ForCausalLM": ("internlm2", "InternLM2ForCausalLM"),
95
    "InternLM2VEForCausalLM": ("internlm2_ve", "InternLM2VEForCausalLM"),
96
    "InternLM3ForCausalLM": ("llama", "LlamaForCausalLM"),
97
98
    "JAISLMHeadModel": ("jais", "JAISLMHeadModel"),
    "JambaForCausalLM": ("jamba", "JambaForCausalLM"),
99
    "Lfm2ForCausalLM": ("lfm2", "Lfm2ForCausalLM"),
100
    "LlamaForCausalLM": ("llama", "LlamaForCausalLM"),
101
    "Llama4ForCausalLM": ("llama4", "Llama4ForCausalLM"),  # noqa: E501
102
103
    # For decapoda-research/llama-*
    "LLaMAForCausalLM": ("llama", "LlamaForCausalLM"),
104
    "MambaForCausalLM": ("mamba", "MambaForCausalLM"),
105
    "FalconMambaForCausalLM": ("mamba", "MambaForCausalLM"),
Dhia Eddine Rhaiem's avatar
Dhia Eddine Rhaiem committed
106
    "FalconH1ForCausalLM":("falcon_h1", "FalconH1ForCausalLM"),
107
    "Mamba2ForCausalLM": ("mamba2", "Mamba2ForCausalLM"),
108
109
    "MiniCPMForCausalLM": ("minicpm", "MiniCPMForCausalLM"),
    "MiniCPM3ForCausalLM": ("minicpm3", "MiniCPM3ForCausalLM"),
110
111
112
113
114
115
    "MistralForCausalLM": ("llama", "LlamaForCausalLM"),
    "MixtralForCausalLM": ("mixtral", "MixtralForCausalLM"),
    "QuantMixtralForCausalLM": ("mixtral_quant", "MixtralForCausalLM"),
    # transformers's mpt class has lower case
    "MptForCausalLM": ("mpt", "MPTForCausalLM"),
    "MPTForCausalLM": ("mpt", "MPTForCausalLM"),
116
    "MiMoForCausalLM": ("mimo", "MiMoForCausalLM"),
117
    "NemotronForCausalLM": ("nemotron", "NemotronForCausalLM"),
Luis Vega's avatar
Luis Vega committed
118
    "NemotronHForCausalLM": ("nemotron_h", "NemotronHForCausalLM"),
119
    "OlmoForCausalLM": ("olmo", "OlmoForCausalLM"),
120
    "Olmo2ForCausalLM": ("olmo2", "Olmo2ForCausalLM"),
121
122
123
124
125
126
127
    "OlmoeForCausalLM": ("olmoe", "OlmoeForCausalLM"),
    "OPTForCausalLM": ("opt", "OPTForCausalLM"),
    "OrionForCausalLM": ("orion", "OrionForCausalLM"),
    "PersimmonForCausalLM": ("persimmon", "PersimmonForCausalLM"),
    "PhiForCausalLM": ("phi", "PhiForCausalLM"),
    "Phi3ForCausalLM": ("phi3", "Phi3ForCausalLM"),
    "PhiMoEForCausalLM": ("phimoe", "PhiMoEForCausalLM"),
128
    "Phi4FlashForCausalLM": ("phi4flash", "Phi4FlashForCausalLM"),
Shinichi Hemmi's avatar
Shinichi Hemmi committed
129
    "Plamo2ForCausalLM": ("plamo2", "Plamo2ForCausalLM"),
130
    "QWenLMHeadModel": ("qwen", "QWenLMHeadModel"),
131
132
    "Qwen2ForCausalLM": ("qwen2", "Qwen2ForCausalLM"),
    "Qwen2MoeForCausalLM": ("qwen2_moe", "Qwen2MoeForCausalLM"),
133
134
    "Qwen3ForCausalLM": ("qwen3", "Qwen3ForCausalLM"),
    "Qwen3MoeForCausalLM": ("qwen3_moe", "Qwen3MoeForCausalLM"),
135
    "RWForCausalLM": ("falcon", "FalconForCausalLM"),
136
    "SeedOssForCausalLM": ("seed_oss", "SeedOssForCausalLM"),
Song's avatar
Song committed
137
    "Step3TextForCausalLM": ("step3_text", "Step3TextForCausalLM"),
138
139
140
141
    "StableLMEpochForCausalLM": ("stablelm", "StablelmForCausalLM"),
    "StableLmForCausalLM": ("stablelm", "StablelmForCausalLM"),
    "Starcoder2ForCausalLM": ("starcoder2", "Starcoder2ForCausalLM"),
    "SolarForCausalLM": ("solar", "SolarForCausalLM"),
142
    "TeleChat2ForCausalLM": ("telechat2", "TeleChat2ForCausalLM"),
143
    "TeleFLMForCausalLM": ("teleflm", "TeleFLMForCausalLM"),
144
    "XverseForCausalLM": ("llama", "LlamaForCausalLM"),
145
    "Zamba2ForCausalLM": ("zamba2", "Zamba2ForCausalLM"),
146
147
148
    # [Encoder-decoder]
    "BartModel": ("bart", "BartForConditionalGeneration"),
    "BartForConditionalGeneration": ("bart", "BartForConditionalGeneration"),
汪志鹏's avatar
汪志鹏 committed
149
    "MBartForConditionalGeneration": ("bart", "MBartForConditionalGeneration"),
150
151
152
}

_EMBEDDING_MODELS = {
153
    # [Text-only]
154
    "BertModel": ("bert", "BertEmbeddingModel"),
155
    "DeciLMForCausalLM": ("nemotron_nas", "DeciLMForCausalLM"),
156
    "Gemma2Model": ("gemma2", "Gemma2ForCausalLM"),
157
    "GlmForCausalLM": ("glm", "GlmForCausalLM"),
158
    "GPT2ForSequenceClassification": ("gpt2", "GPT2ForSequenceClassification"),
159
    "GritLM": ("gritlm", "GritLM"),
160
161
    "GteModel": ("bert_with_rope", "SnowflakeGteNewModel"),
    "GteNewModel": ("bert_with_rope", "GteNewModel"),
162
    "InternLM2ForRewardModel": ("internlm2", "InternLM2ForRewardModel"),
163
    "JambaForSequenceClassification": ("jamba", "JambaForSequenceClassification"),  # noqa: E501
164
    "LlamaModel": ("llama", "LlamaForCausalLM"),
165
166
167
168
169
    **{
        # 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"
    },
170
    "MistralModel": ("llama", "LlamaForCausalLM"),
171
    "ModernBertModel": ("modernbert", "ModernBertModel"),
172
    "NomicBertModel": ("bert_with_rope", "NomicBertModel"),
173
    "Phi3ForCausalLM": ("phi3", "Phi3ForCausalLM"),
174
    "Qwen2Model": ("qwen2", "Qwen2ForCausalLM"),
175
    "Qwen2ForCausalLM": ("qwen2", "Qwen2ForCausalLM"),
176
    "Qwen2ForRewardModel": ("qwen2_rm", "Qwen2ForRewardModel"),
177
    "Qwen2ForProcessRewardModel": ("qwen2_rm", "Qwen2ForProcessRewardModel"),
178
179
    "RobertaForMaskedLM": ("roberta", "RobertaEmbeddingModel"),
    "RobertaModel": ("roberta", "RobertaEmbeddingModel"),
180
    "TeleChat2ForCausalLM": ("telechat2", "TeleChat2ForCausalLM"),
181
    "XLMRobertaModel": ("roberta", "RobertaEmbeddingModel"),
182
    # [Multimodal]
Cyrus Leung's avatar
Cyrus Leung committed
183
    "LlavaNextForConditionalGeneration": ("llava_next", "LlavaNextForConditionalGeneration"),  # noqa: E501
184
    "Phi3VForCausalLM": ("phi3v", "Phi3VForCausalLM"),
185
    "Qwen2VLForConditionalGeneration": ("qwen2_vl", "Qwen2VLForConditionalGeneration"),  # noqa: E501
186
187
188
189
    # 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"),
190
191
}

192
193
_CROSS_ENCODER_MODELS = {
    "BertForSequenceClassification": ("bert", "BertForSequenceClassification"),
194
195
196
197
    "GteNewForSequenceClassification": ("bert_with_rope",
                                        "GteNewForSequenceClassification"),
    "ModernBertForSequenceClassification": ("modernbert",
                                            "ModernBertForSequenceClassification"),
198
199
200
201
    "RobertaForSequenceClassification": ("roberta",
                                         "RobertaForSequenceClassification"),
    "XLMRobertaForSequenceClassification": ("roberta",
                                            "RobertaForSequenceClassification"),
202
    # [Auto-converted (see adapters.py)]
203
    "JinaVLForRanking": ("jina_vl", "JinaVLForSequenceClassification"), # noqa: E501,
204
205
}

206
_MULTIMODAL_MODELS = {
207
    # [Decoder-only]
208
    "AriaForConditionalGeneration": ("aria", "AriaForConditionalGeneration"),
Jennifer Zhao's avatar
Jennifer Zhao committed
209
    "AyaVisionForConditionalGeneration": ("aya_vision", "AyaVisionForConditionalGeneration"),  # noqa: E501
210
211
    "Blip2ForConditionalGeneration": ("blip2", "Blip2ForConditionalGeneration"),
    "ChameleonForConditionalGeneration": ("chameleon", "ChameleonForConditionalGeneration"),  # noqa: E501
212
    "Cohere2VisionForConditionalGeneration": ("cohere2_vision", "Cohere2VisionForConditionalGeneration"),  # noqa: E501
213
    "DeepseekVLV2ForCausalLM": ("deepseek_vl2", "DeepseekVLV2ForCausalLM"),
214
    "Ernie4_5_VLMoeForConditionalGeneration": ("ernie45_vl", "Ernie4_5_VLMoeForConditionalGeneration"),  # noqa: E501
215
    "FuyuForCausalLM": ("fuyu", "FuyuForCausalLM"),
216
    "Gemma3ForConditionalGeneration": ("gemma3_mm", "Gemma3ForConditionalGeneration"),  # noqa: E501
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
217
    "Gemma3nForConditionalGeneration": ("gemma3n_mm", "Gemma3nForConditionalGeneration"),    # noqa: E501
218
    "GLM4VForCausalLM": ("glm4v", "GLM4VForCausalLM"),
219
    "Glm4vForConditionalGeneration": ("glm4_1v", "Glm4vForConditionalGeneration"),  # noqa: E501
Jee Jee Li's avatar
Jee Jee Li committed
220
    "Glm4vMoeForConditionalGeneration": ("glm4_1v", "Glm4vMoeForConditionalGeneration"),  # noqa: E501
221
    "GraniteSpeechForConditionalGeneration": ("granite_speech", "GraniteSpeechForConditionalGeneration"),  # noqa: E501
222
    "H2OVLChatModel": ("h2ovl", "H2OVLChatModel"),
223
    "InternVLChatModel": ("internvl", "InternVLChatModel"),
Lyu Han's avatar
Lyu Han committed
224
    "InternS1ForConditionalGeneration": ("interns1", "InternS1ForConditionalGeneration"),  # noqa: E501
225
    "InternVLForConditionalGeneration": ("interns1", "InternS1ForConditionalGeneration"),  # noqa: E501
226
    "Idefics3ForConditionalGeneration":("idefics3","Idefics3ForConditionalGeneration"),
227
    "SmolVLMForConditionalGeneration": ("smolvlm","SmolVLMForConditionalGeneration"),  # noqa: E501
228
    "KeyeForConditionalGeneration": ("keye", "KeyeForConditionalGeneration"),
229
    "RForConditionalGeneration": ("rvl", "RForConditionalGeneration"),
230
    "KimiVLForConditionalGeneration": ("kimi_vl", "KimiVLForConditionalGeneration"),  # noqa: E501
231
    "Llama_Nemotron_Nano_VL": ("nemotron_vl", "LlamaNemotronVLChatModel"),
232
233
234
235
    "LlavaForConditionalGeneration": ("llava", "LlavaForConditionalGeneration"),
    "LlavaNextForConditionalGeneration": ("llava_next", "LlavaNextForConditionalGeneration"),  # noqa: E501
    "LlavaNextVideoForConditionalGeneration": ("llava_next_video", "LlavaNextVideoForConditionalGeneration"),  # noqa: E501
    "LlavaOnevisionForConditionalGeneration": ("llava_onevision", "LlavaOnevisionForConditionalGeneration"),  # noqa: E501
236
    "MantisForConditionalGeneration": ("llava", "MantisForConditionalGeneration"),  # noqa: E501
237
    "MiniMaxVL01ForConditionalGeneration": ("minimax_vl_01", "MiniMaxVL01ForConditionalGeneration"),  # noqa: E501
238
    "MiniCPMO": ("minicpmo", "MiniCPMO"),
239
    "MiniCPMV": ("minicpmv", "MiniCPMV"),
240
    "Mistral3ForConditionalGeneration": ("mistral3", "Mistral3ForConditionalGeneration"),  # noqa: E501
241
    "MolmoForCausalLM": ("molmo", "MolmoForCausalLM"),
242
    "NVLM_D": ("nvlm_d", "NVLM_D_Model"),
243
    "Ovis": ("ovis", "Ovis"),
244
    "Ovis2_5": ("ovis2_5", "Ovis2_5"),
245
    "PaliGemmaForConditionalGeneration": ("paligemma", "PaliGemmaForConditionalGeneration"),  # noqa: E501
246
    "Phi3VForCausalLM": ("phi3v", "Phi3VForCausalLM"),
247
248
    "Phi4MMForCausalLM": ("phi4mm", "Phi4MMForCausalLM"),
    "Phi4MultimodalForCausalLM": ("phi4_multimodal", "Phi4MultimodalForCausalLM"),  # noqa: E501
249
    "PixtralForConditionalGeneration": ("pixtral", "PixtralForConditionalGeneration"),  # noqa: E501
250
    "QwenVLForConditionalGeneration": ("qwen_vl", "QwenVLForConditionalGeneration"),  # noqa: E501
251
    "Qwen2VLForConditionalGeneration": ("qwen2_vl", "Qwen2VLForConditionalGeneration"),  # noqa: E501
Roger Wang's avatar
Roger Wang committed
252
    "Qwen2_5_VLForConditionalGeneration": ("qwen2_5_vl", "Qwen2_5_VLForConditionalGeneration"),  # noqa: E501
253
    "Qwen2AudioForConditionalGeneration": ("qwen2_audio", "Qwen2AudioForConditionalGeneration"),  # noqa: E501
254
    "Qwen2_5OmniModel": ("qwen2_5_omni_thinker", "Qwen2_5OmniThinkerForConditionalGeneration"),  # noqa: E501
255
    "Qwen2_5OmniForConditionalGeneration": ("qwen2_5_omni_thinker", "Qwen2_5OmniThinkerForConditionalGeneration"),  # noqa: E501
256
    "UltravoxModel": ("ultravox", "UltravoxModel"),
Song's avatar
Song committed
257
    "Step3VLForConditionalGeneration": ("step3_vl", "Step3VLForConditionalGeneration"),  # noqa: E501
汪志鹏's avatar
汪志鹏 committed
258
    "TarsierForConditionalGeneration": ("tarsier", "TarsierForConditionalGeneration"),  # noqa: E501
259
    "Tarsier2ForConditionalGeneration": ("qwen2_vl", "Tarsier2ForConditionalGeneration"),  # noqa: E501
Patrick von Platen's avatar
Patrick von Platen committed
260
    "VoxtralForConditionalGeneration": ("voxtral", "VoxtralForConditionalGeneration"),  # noqa: E501
261
    # [Encoder-decoder]
汪志鹏's avatar
汪志鹏 committed
262
    "DonutForConditionalGeneration": ("donut", "DonutForConditionalGeneration"),
263
    "Florence2ForConditionalGeneration": ("florence2", "Florence2ForConditionalGeneration"),  # noqa: E501
264
    "MllamaForConditionalGeneration": ("mllama", "MllamaForConditionalGeneration"),  # noqa: E501
265
    "Llama4ForConditionalGeneration": ("mllama4", "Llama4ForConditionalGeneration"),  # noqa: E501
266
    "SkyworkR1VChatModel": ("skyworkr1v", "SkyworkR1VChatModel"),
267
    "WhisperForConditionalGeneration": ("whisper", "WhisperForConditionalGeneration"),  # noqa: E501
268
}
269
270

_SPECULATIVE_DECODING_MODELS = {
271
    "MiMoMTPModel": ("mimo_mtp", "MiMoMTP"),
272
    "EagleLlamaForCausalLM": ("llama_eagle", "EagleLlamaForCausalLM"),
zhiweiz's avatar
zhiweiz committed
273
    "EagleLlama4ForCausalLM": ("llama4_eagle", "EagleLlama4ForCausalLM"),
274
    "EagleMiniCPMForCausalLM": ("minicpm_eagle", "EagleMiniCPMForCausalLM"),
275
    "Eagle3LlamaForCausalLM": ("llama_eagle3", "Eagle3LlamaForCausalLM"),
276
277
    # TODO: Re-enable this once tests/models/test_initialization.py is fixed, see PR #22333 #22611  # noqa: E501
    # "LlamaForCausalLMEagle3": ("llama_eagle3", "Eagle3LlamaForCausalLM"),
278
    "EagleDeepSeekMTPModel": ("deepseek_eagle", "EagleDeepseekV3ForCausalLM"),
279
    "DeepSeekMTPModel": ("deepseek_mtp", "DeepSeekMTP"),
280
    "ErnieMTPModel": ("ernie_mtp", "ErnieMTP"),
Yuxuan Zhang's avatar
Yuxuan Zhang committed
281
    "Glm4MoeMTPModel": ("glm4_moe_mtp", "Glm4MoeMTP"),
282
    "MedusaModel": ("medusa", "Medusa"),
283
284
285
    # Temporarily disabled.
    # # TODO(woosuk): Re-enable this once the MLP Speculator is supported in V1.
    # "MLPSpeculatorPreTrainedModel": ("mlp_speculator", "MLPSpeculator"),
286
}
287

288
_TRANSFORMERS_SUPPORTED_MODELS = {
289
290
291
    # Text generation models
    "SmolLM3ForCausalLM": ("transformers", "TransformersForCausalLM"),
    # Multimodal models
292
293
294
295
    "Emu3ForConditionalGeneration": ("transformers", "TransformersForMultimodalLM"),  # noqa: E501
}

_TRANSFORMERS_BACKEND_MODELS = {
296
    "TransformersModel": ("transformers", "TransformersModel"),
297
    "TransformersForCausalLM": ("transformers", "TransformersForCausalLM"),
298
    "TransformersForMultimodalLM": ("transformers", "TransformersForMultimodalLM"), # noqa: E501
299
}
300
# yapf: enable
301

302
_VLLM_MODELS = {
303
    **_TEXT_GENERATION_MODELS,
304
    **_EMBEDDING_MODELS,
305
    **_CROSS_ENCODER_MODELS,
306
    **_MULTIMODAL_MODELS,
307
    **_SPECULATIVE_DECODING_MODELS,
308
309
    **_TRANSFORMERS_SUPPORTED_MODELS,
    **_TRANSFORMERS_BACKEND_MODELS,
310
311
}

312
313
314
315
316
317
318
319
# 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"
]

320
321
_PREVIOUSLY_SUPPORTED_MODELS = {"Phi3SmallForCausalLM": "0.9.2"}

322

323
324
@dataclass(frozen=True)
class _ModelInfo:
325
    architecture: str
326
    is_text_generation_model: bool
327
    is_pooling_model: bool
328
    default_pooling_type: str
329
    supports_cross_encoding: bool
330
    supports_multimodal: bool
331
    supports_multimodal_raw_input_only: bool
332
    supports_multimodal_encoder_tp_data: bool
333
    supports_pp: bool
334
335
    has_inner_state: bool
    is_attention_free: bool
336
    is_hybrid: bool
337
    has_noops: bool
338
    supports_transcription: bool
339
    supports_transcription_only: bool
340
    supports_v0_only: bool
341
342

    @staticmethod
343
    def from_model_cls(model: type[nn.Module]) -> "_ModelInfo":
344
        return _ModelInfo(
345
            architecture=model.__name__,
346
            is_text_generation_model=is_text_generation_model(model),
347
            is_pooling_model=is_pooling_model(model),
348
            default_pooling_type=get_default_pooling_type(model),
349
            supports_cross_encoding=supports_cross_encoding(model),
350
            supports_multimodal=supports_multimodal(model),
351
352
            supports_multimodal_raw_input_only=
            supports_multimodal_raw_input_only(model),
353
354
            supports_multimodal_encoder_tp_data=
            supports_multimodal_encoder_tp_data(model),
355
            supports_pp=supports_pp(model),
356
357
            has_inner_state=has_inner_state(model),
            is_attention_free=is_attention_free(model),
358
            is_hybrid=is_hybrid(model),
359
            supports_transcription=supports_transcription(model),
360
361
            supports_transcription_only=(supports_transcription(model) and
                                         model.supports_transcription_only),
362
            supports_v0_only=supports_v0_only(model),
363
            has_noops=has_noops(model),
364
        )
365
366


367
class _BaseRegisteredModel(ABC):
368

369
370
371
    @abstractmethod
    def inspect_model_cls(self) -> _ModelInfo:
        raise NotImplementedError
372

373
    @abstractmethod
374
    def load_model_cls(self) -> type[nn.Module]:
375
        raise NotImplementedError
376
377


378
379
380
381
382
383
384
@dataclass(frozen=True)
class _RegisteredModel(_BaseRegisteredModel):
    """
    Represents a model that has already been imported in the main process.
    """

    interfaces: _ModelInfo
385
    model_cls: type[nn.Module]
386
387

    @staticmethod
388
    def from_model_cls(model_cls: type[nn.Module]):
389
390
391
392
393
394
395
396
        return _RegisteredModel(
            interfaces=_ModelInfo.from_model_cls(model_cls),
            model_cls=model_cls,
        )

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

397
    def load_model_cls(self) -> type[nn.Module]:
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
        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()))

414
    def load_model_cls(self) -> type[nn.Module]:
415
416
417
418
419
420
421
422
        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,
423
) -> Optional[type[nn.Module]]:
424
    from vllm.platforms import current_platform
425
    current_platform.verify_model_arch(model_arch)
426
427
428
429
430
431
    try:
        return model.load_model_cls()
    except Exception:
        logger.exception("Error in loading model architecture '%s'",
                         model_arch)
        return None
432
433


434
435
436
437
438
439
440
441
442
443
444
@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
445
446


447
448
449
@dataclass
class _ModelRegistry:
    # Keyed by model_arch
450
    models: dict[str, _BaseRegisteredModel] = field(default_factory=dict)
451

452
    def get_supported_archs(self) -> Set[str]:
453
        return self.models.keys()
454

455
456
457
    def register_model(
        self,
        model_arch: str,
458
        model_cls: Union[type[nn.Module], str],
459
    ) -> None:
460
461
462
        """
        Register an external model to be used in vLLM.

463
        `model_cls` can be either:
464

465
        - A [`torch.nn.Module`][] class directly referencing the model.
466
        - A string in the format `<module>:<class>` which can be used to
467
468
          lazily import the model. This is useful to avoid initializing CUDA
          when importing the model and thus the related error
469
          `RuntimeError: Cannot re-initialize CUDA in forked subprocess`.
470
        """
471
472
473
474
        if not isinstance(model_arch, str):
            msg = f"`model_arch` should be a string, not a {type(model_arch)}"
            raise TypeError(msg)

475
        if model_arch in self.models:
476
477
478
            logger.warning(
                "Model architecture %s is already registered, and will be "
                "overwritten by the new model class %s.", model_arch,
479
480
481
482
483
484
485
                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)
486

487
            model = _LazyRegisteredModel(*split_str)
488
        elif isinstance(model_cls, type) and issubclass(model_cls, nn.Module):
489
            model = _RegisteredModel.from_model_cls(model_cls)
490
491
492
493
        else:
            msg = ("`model_cls` should be a string or PyTorch model class, "
                   f"not a {type(model_arch)}")
            raise TypeError(msg)
494

495
        self.models[model_arch] = model
496

497
    def _raise_for_unsupported(self, architectures: list[str]):
498
        all_supported_archs = self.get_supported_archs()
499

500
501
502
503
504
        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.")

505
506
507
508
509
510
511
512
513
514
        for arch in architectures:
            if arch in _PREVIOUSLY_SUPPORTED_MODELS:
                previous_version = _PREVIOUSLY_SUPPORTED_MODELS[arch]

                raise ValueError(
                    f"Model architecture {arch} was supported in vLLM until "
                    f"v{previous_version}, and is not supported anymore. "
                    "Please use an older version of vLLM if you want to "
                    "use this model architecture.")

515
516
517
        raise ValueError(
            f"Model architectures {architectures} are not supported for now. "
            f"Supported architectures: {all_supported_archs}")
518

519
    def _try_load_model_cls(self,
520
                            model_arch: str) -> Optional[type[nn.Module]]:
521
522
        if model_arch not in self.models:
            return None
523

524
        return _try_load_model_cls(model_arch, self.models[model_arch])
525

526
    def _try_inspect_model_cls(self, model_arch: str) -> Optional[_ModelInfo]:
527
528
        if model_arch not in self.models:
            return None
529

530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
        return _try_inspect_model_cls(model_arch, self.models[model_arch])

    def _try_resolve_transformers(
        self,
        architecture: str,
        model_config: ModelConfig,
    ) -> Optional[str]:
        if architecture in _TRANSFORMERS_BACKEND_MODELS:
            return architecture

        auto_map: dict[str, str] = getattr(model_config.hf_config, "auto_map",
                                           None) or dict()

        # Make sure that config class is always initialized before model class,
        # otherwise the model class won't be able to access the config class,
        # the expected auto_map should have correct order like:
        # "auto_map": {
        #     "AutoConfig": "<your-repo-name>--<config-name>",
        #     "AutoModel": "<your-repo-name>--<config-name>",
        #     "AutoModelFor<Task>": "<your-repo-name>--<config-name>",
        # },
        for prefix in ("AutoConfig", "AutoModel"):
            for name, module in auto_map.items():
                if name.startswith(prefix):
                    try_get_class_from_dynamic_module(
                        module,
                        model_config.model,
                        revision=model_config.revision,
                        warn_on_fail=False,
                    )

        model_module = getattr(transformers, architecture, None)

        if model_module is None:
            for name, module in auto_map.items():
                if name.startswith("AutoModel"):
                    model_module = try_get_class_from_dynamic_module(
                        module,
                        model_config.model,
                        revision=model_config.revision,
                        warn_on_fail=True,
                    )
                    if model_module is not None:
                        break
            else:
                if model_config.model_impl != ModelImpl.TRANSFORMERS:
                    return None

                raise ValueError(
                    f"Cannot find model module. {architecture!r} is not a "
                    "registered model in the Transformers library (only "
                    "relevant if the model is meant to be in Transformers) "
                    "and 'AutoModel' is not present in the model config's "
                    "'auto_map' (relevant if the model is custom).")

        if not model_module.is_backend_compatible():
            if model_config.model_impl != ModelImpl.TRANSFORMERS:
587
                return None
588

589
590
591
            raise ValueError(
                f"The Transformers implementation of {architecture!r} "
                "is not compatible with vLLM.")
592

593
        return model_config._get_transformers_backend_cls()
594

595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
    def _normalize_arch(
        self,
        architecture: str,
        model_config: ModelConfig,
    ) -> str:
        if architecture in self.models:
            return architecture

        # This may be called in order to resolve runner_type and convert_type
        # in the first place, in which case we consider the default match
        match = try_match_architecture_defaults(
            architecture,
            runner_type=getattr(model_config, "runner_type", None),
            convert_type=getattr(model_config, "convert_type", None),
        )
        if match:
            suffix, _ = match

            # Get the name of the base model to convert
            for repl_suffix, _ in iter_architecture_defaults():
                base_arch = architecture.replace(suffix, repl_suffix)
                if base_arch in self.models:
                    return base_arch

        return architecture
620

621
622
    def inspect_model_cls(
        self,
623
        architectures: Union[str, list[str]],
624
        model_config: ModelConfig,
625
    ) -> tuple[_ModelInfo, str]:
626
627
        if isinstance(architectures, str):
            architectures = [architectures]
628
629
        if not architectures:
            raise ValueError("No model architectures are specified")
630
631
632
633
634
635
636
637
638
639

        # Require transformers impl
        if model_config.model_impl == ModelImpl.TRANSFORMERS:
            arch = self._try_resolve_transformers(architectures[0],
                                                  model_config)
            if arch is not None:
                model_info = self._try_inspect_model_cls(arch)
                if model_info is not None:
                    return (model_info, arch)

640
641
642
643
644
645
646
647
648
649
650
651
652
        # Fallback to transformers impl (after resolving convert_type)
        if (all(arch not in self.models for arch in architectures)
                and model_config.model_impl == ModelImpl.AUTO
                and getattr(model_config, "convert_type", "none") == "none"):
            arch = self._try_resolve_transformers(architectures[0],
                                                  model_config)
            if arch is not None:
                model_info = self._try_inspect_model_cls(arch)
                if model_info is not None:
                    return (model_info, arch)

        for arch in architectures:
            normalized_arch = self._normalize_arch(arch, model_config)
653
            model_info = self._try_inspect_model_cls(normalized_arch)
654
            if model_info is not None:
655
                return (model_info, arch)
656

657
658
659
        # Fallback to transformers impl (before resolving runner_type)
        if (all(arch not in self.models for arch in architectures)
                and model_config.model_impl == ModelImpl.AUTO):
660
661
662
663
664
665
666
            arch = self._try_resolve_transformers(architectures[0],
                                                  model_config)
            if arch is not None:
                model_info = self._try_inspect_model_cls(arch)
                if model_info is not None:
                    return (model_info, arch)

667
        return self._raise_for_unsupported(architectures)
668

669
670
    def resolve_model_cls(
        self,
671
        architectures: Union[str, list[str]],
672
        model_config: ModelConfig,
673
    ) -> tuple[type[nn.Module], str]:
674
675
        if isinstance(architectures, str):
            architectures = [architectures]
676
677
        if not architectures:
            raise ValueError("No model architectures are specified")
678
679
680
681
682
683
684
685
686
687

        # Require transformers impl
        if model_config.model_impl == ModelImpl.TRANSFORMERS:
            arch = self._try_resolve_transformers(architectures[0],
                                                  model_config)
            if arch is not None:
                model_cls = self._try_load_model_cls(arch)
                if model_cls is not None:
                    return (model_cls, arch)

688
689
690
691
692
693
694
695
696
697
698
699
700
        # Fallback to transformers impl (after resolving convert_type)
        if (all(arch not in self.models for arch in architectures)
                and model_config.model_impl == ModelImpl.AUTO
                and getattr(model_config, "convert_type", "none") == "none"):
            arch = self._try_resolve_transformers(architectures[0],
                                                  model_config)
            if arch is not None:
                model_cls = self._try_load_model_cls(arch)
                if model_cls is not None:
                    return (model_cls, arch)

        for arch in architectures:
            normalized_arch = self._normalize_arch(arch, model_config)
701
            model_cls = self._try_load_model_cls(normalized_arch)
702
703
            if model_cls is not None:
                return (model_cls, arch)
704

705
706
707
        # Fallback to transformers impl (before resolving runner_type)
        if (all(arch not in self.models for arch in architectures)
                and model_config.model_impl == ModelImpl.AUTO):
708
709
710
711
712
713
714
            arch = self._try_resolve_transformers(architectures[0],
                                                  model_config)
            if arch is not None:
                model_cls = self._try_load_model_cls(arch)
                if model_cls is not None:
                    return (model_cls, arch)

715
        return self._raise_for_unsupported(architectures)
716

717
718
    def is_text_generation_model(
        self,
719
        architectures: Union[str, list[str]],
720
        model_config: ModelConfig,
721
    ) -> bool:
722
        model_cls, _ = self.inspect_model_cls(architectures, model_config)
723
        return model_cls.is_text_generation_model
724

725
    def is_pooling_model(
726
        self,
727
        architectures: Union[str, list[str]],
728
        model_config: ModelConfig,
729
    ) -> bool:
730
        model_cls, _ = self.inspect_model_cls(architectures, model_config)
731
        return model_cls.is_pooling_model
732

733
734
    def is_cross_encoder_model(
        self,
735
        architectures: Union[str, list[str]],
736
        model_config: ModelConfig,
737
    ) -> bool:
738
        model_cls, _ = self.inspect_model_cls(architectures, model_config)
739
        return model_cls.supports_cross_encoding
740

741
742
    def is_multimodal_model(
        self,
743
        architectures: Union[str, list[str]],
744
        model_config: ModelConfig,
745
    ) -> bool:
746
        model_cls, _ = self.inspect_model_cls(architectures, model_config)
747
        return model_cls.supports_multimodal
748

749
    def is_multimodal_raw_input_only_model(
750
751
        self,
        architectures: Union[str, list[str]],
752
        model_config: ModelConfig,
753
    ) -> bool:
754
        model_cls, _ = self.inspect_model_cls(architectures, model_config)
755
        return model_cls.supports_multimodal_raw_input_only
756

757
758
    def is_pp_supported_model(
        self,
759
        architectures: Union[str, list[str]],
760
        model_config: ModelConfig,
761
    ) -> bool:
762
        model_cls, _ = self.inspect_model_cls(architectures, model_config)
763
        return model_cls.supports_pp
764

765
766
    def model_has_inner_state(
        self,
767
        architectures: Union[str, list[str]],
768
        model_config: ModelConfig,
769
    ) -> bool:
770
        model_cls, _ = self.inspect_model_cls(architectures, model_config)
771
        return model_cls.has_inner_state
772

773
774
    def is_attention_free_model(
        self,
775
        architectures: Union[str, list[str]],
776
        model_config: ModelConfig,
777
    ) -> bool:
778
        model_cls, _ = self.inspect_model_cls(architectures, model_config)
779
        return model_cls.is_attention_free
780

781
782
    def is_hybrid_model(
        self,
783
        architectures: Union[str, list[str]],
784
        model_config: ModelConfig,
785
    ) -> bool:
786
        model_cls, _ = self.inspect_model_cls(architectures, model_config)
787
788
        return model_cls.is_hybrid

789
790
    def is_noops_model(
        self,
791
        architectures: Union[str, list[str]],
792
        model_config: ModelConfig,
793
    ) -> bool:
794
        model_cls, _ = self.inspect_model_cls(architectures, model_config)
795
796
        return model_cls.has_noops

797
798
    def is_transcription_model(
        self,
799
        architectures: Union[str, list[str]],
800
        model_config: ModelConfig,
801
    ) -> bool:
802
        model_cls, _ = self.inspect_model_cls(architectures, model_config)
803
804
        return model_cls.supports_transcription

805
806
807
    def is_transcription_only_model(
        self,
        architectures: Union[str, list[str]],
808
        model_config: ModelConfig,
809
    ) -> bool:
810
        model_cls, _ = self.inspect_model_cls(architectures, model_config)
811
812
        return model_cls.supports_transcription_only

813
814
    def is_v1_compatible(
        self,
815
        architectures: Union[str, list[str]],
816
        model_config: ModelConfig,
817
    ) -> bool:
818
        model_cls, _ = self.inspect_model_cls(architectures, model_config)
819
820
        return not model_cls.supports_v0_only

821
822

ModelRegistry = _ModelRegistry({
823
824
    model_arch:
    _LazyRegisteredModel(
825
826
827
828
829
830
831
832
833
834
        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:
835
836
837
838
839
    # 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")

840
        # `cloudpickle` allows pickling lambda functions directly
841
        import cloudpickle
842
        input_bytes = cloudpickle.dumps((fn, output_filepath))
843
844
845

        # cannot use `sys.executable __file__` here because the script
        # contains relative imports
846
847
848
        returned = subprocess.run(_SUBPROCESS_COMMAND,
                                  input=input_bytes,
                                  capture_output=True)
849
850
851
852
853
854
855
856
857

        # 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

858
        with open(output_filepath, "rb") as f:
859
860
861
862
863
864
865
866
867
868
869
            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()
870
871
872

    with open(output_file, "wb") as f:
        f.write(pickle.dumps(result))
873
874
875


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