registry.py 38.5 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
    "ApertusForCausalLM": ("apertus", "ApertusForCausalLM"),
43
44
    "AquilaModel": ("llama", "LlamaForCausalLM"),
    "AquilaForCausalLM": ("llama", "LlamaForCausalLM"),  # AquilaChat2
Raghav Ravishankar's avatar
Raghav Ravishankar committed
45
    "ArceeForCausalLM": ("arcee", "ArceeForCausalLM"),
46
    "ArcticForCausalLM": ("arctic", "ArcticForCausalLM"),
47
    "MiniMaxForCausalLM": ("minimax_text_01", "MiniMaxText01ForCausalLM"),
48
    "MiniMaxText01ForCausalLM": ("minimax_text_01", "MiniMaxText01ForCausalLM"),
49
    "MiniMaxM1ForCausalLM": ("minimax_text_01", "MiniMaxText01ForCausalLM"),
50
51
52
53
    # baichuan-7b, upper case 'C' in the class name
    "BaiChuanForCausalLM": ("baichuan", "BaiChuanForCausalLM"),
    # baichuan-13b, lower case 'c' in the class name
    "BaichuanForCausalLM": ("baichuan", "BaichuanForCausalLM"),
54
    "BailingMoeForCausalLM": ("bailing_moe", "BailingMoeForCausalLM"),
Yu Chin Fabian Lim's avatar
Yu Chin Fabian Lim committed
55
    "BambaForCausalLM": ("bamba", "BambaForCausalLM"),
56
    "BloomForCausalLM": ("bloom", "BloomForCausalLM"),
57
    "ChatGLMModel": ("chatglm", "ChatGLMForCausalLM"),
58
    "ChatGLMForConditionalGeneration": ("chatglm", "ChatGLMForCausalLM"),
59
    "CohereForCausalLM": ("commandr", "CohereForCausalLM"),
60
    "Cohere2ForCausalLM": ("commandr", "CohereForCausalLM"),
61
    "DbrxForCausalLM": ("dbrx", "DbrxForCausalLM"),
62
    "DeciLMForCausalLM": ("nemotron_nas", "DeciLMForCausalLM"),
63
64
    "DeepseekForCausalLM": ("deepseek", "DeepseekForCausalLM"),
    "DeepseekV2ForCausalLM": ("deepseek_v2", "DeepseekV2ForCausalLM"),
65
    "DeepseekV3ForCausalLM": ("deepseek_v2", "DeepseekV3ForCausalLM"),
66
    "Dots1ForCausalLM": ("dots1", "Dots1ForCausalLM"),
67
    "Ernie4_5ForCausalLM": ("ernie45", "Ernie4_5ForCausalLM"),
68
    "Ernie4_5_MoeForCausalLM": ("ernie45_moe", "Ernie4_5_MoeForCausalLM"),
69
    "ExaoneForCausalLM": ("exaone", "ExaoneForCausalLM"),
70
    "Exaone4ForCausalLM": ("exaone4", "Exaone4ForCausalLM"),
71
    "FalconForCausalLM": ("falcon", "FalconForCausalLM"),
72
    "Fairseq2LlamaForCausalLM": ("fairseq2_llama", "Fairseq2LlamaForCausalLM"),
73
74
    "GemmaForCausalLM": ("gemma", "GemmaForCausalLM"),
    "Gemma2ForCausalLM": ("gemma2", "Gemma2ForCausalLM"),
75
    "Gemma3ForCausalLM": ("gemma3", "Gemma3ForCausalLM"),
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
76
    "Gemma3nForCausalLM": ("gemma3n", "Gemma3nForCausalLM"),
77
    "GlmForCausalLM": ("glm", "GlmForCausalLM"),
Yuxuan Zhang's avatar
Yuxuan Zhang committed
78
    "Glm4ForCausalLM": ("glm4", "Glm4ForCausalLM"),
Yuxuan Zhang's avatar
Yuxuan Zhang committed
79
    "Glm4MoeForCausalLM": ("glm4_moe", "Glm4MoeForCausalLM"),
80
    "GptOssForCausalLM": ("gpt_oss", "GptOssForCausalLM"),
81
82
83
84
85
86
    "GPT2LMHeadModel": ("gpt2", "GPT2LMHeadModel"),
    "GPTBigCodeForCausalLM": ("gpt_bigcode", "GPTBigCodeForCausalLM"),
    "GPTJForCausalLM": ("gpt_j", "GPTJForCausalLM"),
    "GPTNeoXForCausalLM": ("gpt_neox", "GPTNeoXForCausalLM"),
    "GraniteForCausalLM": ("granite", "GraniteForCausalLM"),
    "GraniteMoeForCausalLM": ("granitemoe", "GraniteMoeForCausalLM"),
87
    "GraniteMoeHybridForCausalLM": ("granitemoehybrid", "GraniteMoeHybridForCausalLM"),   # noqa: E501
88
    "GraniteMoeSharedForCausalLM": ("granitemoeshared", "GraniteMoeSharedForCausalLM"),   # noqa: E501
89
    "GritLM": ("gritlm", "GritLM"),
Michael Goin's avatar
Michael Goin committed
90
    "Grok1ModelForCausalLM": ("grok1", "Grok1ForCausalLM"),
91
92
    "HunYuanMoEV1ForCausalLM": ("hunyuan_v1", "HunYuanMoEV1ForCausalLM"),
    "HunYuanDenseV1ForCausalLM": ("hunyuan_v1", "HunYuanDenseV1ForCausalLM"),
93
    "HCXVisionForCausalLM": ("hyperclovax_vision", "HCXVisionForCausalLM"),
94
95
    "InternLMForCausalLM": ("llama", "LlamaForCausalLM"),
    "InternLM2ForCausalLM": ("internlm2", "InternLM2ForCausalLM"),
96
    "InternLM2VEForCausalLM": ("internlm2_ve", "InternLM2VEForCausalLM"),
97
    "InternLM3ForCausalLM": ("llama", "LlamaForCausalLM"),
98
99
    "JAISLMHeadModel": ("jais", "JAISLMHeadModel"),
    "JambaForCausalLM": ("jamba", "JambaForCausalLM"),
100
    "Lfm2ForCausalLM": ("lfm2", "Lfm2ForCausalLM"),
101
    "LlamaForCausalLM": ("llama", "LlamaForCausalLM"),
102
    "Llama4ForCausalLM": ("llama4", "Llama4ForCausalLM"),  # noqa: E501
103
104
    # For decapoda-research/llama-*
    "LLaMAForCausalLM": ("llama", "LlamaForCausalLM"),
105
    "MambaForCausalLM": ("mamba", "MambaForCausalLM"),
106
    "FalconMambaForCausalLM": ("mamba", "MambaForCausalLM"),
Dhia Eddine Rhaiem's avatar
Dhia Eddine Rhaiem committed
107
    "FalconH1ForCausalLM":("falcon_h1", "FalconH1ForCausalLM"),
108
    "Mamba2ForCausalLM": ("mamba2", "Mamba2ForCausalLM"),
109
110
    "MiniCPMForCausalLM": ("minicpm", "MiniCPMForCausalLM"),
    "MiniCPM3ForCausalLM": ("minicpm3", "MiniCPM3ForCausalLM"),
111
112
113
114
115
    "MistralForCausalLM": ("llama", "LlamaForCausalLM"),
    "MixtralForCausalLM": ("mixtral", "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
    "Gemma3TextModel": ("gemma3", "Gemma3Model"),
158
    "GlmForCausalLM": ("glm", "GlmForCausalLM"),
159
    "GPT2ForSequenceClassification": ("gpt2", "GPT2ForSequenceClassification"),
160
    "GritLM": ("gritlm", "GritLM"),
161
162
    "GteModel": ("bert_with_rope", "SnowflakeGteNewModel"),
    "GteNewModel": ("bert_with_rope", "GteNewModel"),
163
    "InternLM2ForRewardModel": ("internlm2", "InternLM2ForRewardModel"),
164
    "JambaForSequenceClassification": ("jamba", "JambaForSequenceClassification"),  # noqa: E501
165
    "LlamaModel": ("llama", "LlamaForCausalLM"),
166
167
168
169
170
    **{
        # 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"
    },
171
    "MistralModel": ("llama", "LlamaForCausalLM"),
172
    "ModernBertModel": ("modernbert", "ModernBertModel"),
173
    "NomicBertModel": ("bert_with_rope", "NomicBertModel"),
174
    "Phi3ForCausalLM": ("phi3", "Phi3ForCausalLM"),
175
    "Qwen2Model": ("qwen2", "Qwen2ForCausalLM"),
176
    "Qwen2ForCausalLM": ("qwen2", "Qwen2ForCausalLM"),
177
    "Qwen2ForRewardModel": ("qwen2_rm", "Qwen2ForRewardModel"),
178
    "Qwen2ForProcessRewardModel": ("qwen2_rm", "Qwen2ForProcessRewardModel"),
179
180
    "RobertaForMaskedLM": ("roberta", "RobertaEmbeddingModel"),
    "RobertaModel": ("roberta", "RobertaEmbeddingModel"),
181
    "TeleChat2ForCausalLM": ("telechat2", "TeleChat2ForCausalLM"),
182
    "XLMRobertaModel": ("roberta", "RobertaEmbeddingModel"),
183
    # [Multimodal]
Cyrus Leung's avatar
Cyrus Leung committed
184
    "LlavaNextForConditionalGeneration": ("llava_next", "LlavaNextForConditionalGeneration"),  # noqa: E501
185
    "Phi3VForCausalLM": ("phi3v", "Phi3VForCausalLM"),
186
    "Qwen2VLForConditionalGeneration": ("qwen2_vl", "Qwen2VLForConditionalGeneration"),  # noqa: E501
187
188
    # Technically Terratorch models work on images, both in
    # input and output. I am adding it here because it piggy-backs on embedding
189
    # models for the time being.
190
191
    "PrithviGeoSpatialMAE": ("terratorch", "Terratorch"),
    "Terratorch": ("terratorch", "Terratorch"),
192
193
}

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

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

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

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

_TRANSFORMERS_BACKEND_MODELS = {
299
    "TransformersModel": ("transformers", "TransformersModel"),
300
    "TransformersForCausalLM": ("transformers", "TransformersForCausalLM"),
301
    "TransformersForMultimodalLM": ("transformers", "TransformersForMultimodalLM"), # noqa: E501
302
}
303
# yapf: enable
304

305
_VLLM_MODELS = {
306
    **_TEXT_GENERATION_MODELS,
307
    **_EMBEDDING_MODELS,
308
    **_CROSS_ENCODER_MODELS,
309
    **_MULTIMODAL_MODELS,
310
    **_SPECULATIVE_DECODING_MODELS,
311
312
    **_TRANSFORMERS_SUPPORTED_MODELS,
    **_TRANSFORMERS_BACKEND_MODELS,
313
314
}

315
316
317
318
319
320
321
322
# 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"
]

323
324
_PREVIOUSLY_SUPPORTED_MODELS = {"Phi3SmallForCausalLM": "0.9.2"}

325

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

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


370
class _BaseRegisteredModel(ABC):
371

372
373
374
    @abstractmethod
    def inspect_model_cls(self) -> _ModelInfo:
        raise NotImplementedError
375

376
    @abstractmethod
377
    def load_model_cls(self) -> type[nn.Module]:
378
        raise NotImplementedError
379
380


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

    interfaces: _ModelInfo
388
    model_cls: type[nn.Module]
389
390

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

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

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

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


437
438
439
440
441
442
443
444
445
446
447
@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
448
449


450
451
452
@dataclass
class _ModelRegistry:
    # Keyed by model_arch
453
    models: dict[str, _BaseRegisteredModel] = field(default_factory=dict)
454

455
    def get_supported_archs(self) -> Set[str]:
456
        return self.models.keys()
457

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

466
        `model_cls` can be either:
467

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

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

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

498
        self.models[model_arch] = model
499

500
    def _raise_for_unsupported(self, architectures: list[str]):
501
        all_supported_archs = self.get_supported_archs()
502

503
504
505
506
507
        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.")

508
509
510
511
512
513
514
515
516
517
        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.")

518
519
520
        raise ValueError(
            f"Model architectures {architectures} are not supported for now. "
            f"Supported architectures: {all_supported_archs}")
521

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

527
        return _try_load_model_cls(model_arch, self.models[model_arch])
528

529
    def _try_inspect_model_cls(self, model_arch: str) -> Optional[_ModelInfo]:
530
531
        if model_arch not in self.models:
            return None
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
587
588
589
        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:
590
                return None
591

592
593
594
            raise ValueError(
                f"The Transformers implementation of {architecture!r} "
                "is not compatible with vLLM.")
595

596
        return model_config._get_transformers_backend_cls()
597

598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
    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
623

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

        # 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)
642
643
644
        elif model_config.model_impl == ModelImpl.TERRATORCH:
            model_info = self._try_inspect_model_cls("Terratorch")
            return (model_info, "Terratorch")
645

646
647
648
649
650
651
652
653
654
655
656
657
658
        # 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)
659
            model_info = self._try_inspect_model_cls(normalized_arch)
660
            if model_info is not None:
661
                return (model_info, arch)
662

663
664
665
        # 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):
666
667
668
669
670
671
672
            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)

673
        return self._raise_for_unsupported(architectures)
674

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

        # 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)
693
694
695
696
697
        elif model_config.model_impl == ModelImpl.TERRATORCH:
            arch = "Terratorch"
            model_cls = self._try_load_model_cls(arch)
            if model_cls is not None:
                return (model_cls, arch)
698

699
700
701
702
703
704
705
706
707
708
709
710
711
        # 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)
712
            model_cls = self._try_load_model_cls(normalized_arch)
713
714
            if model_cls is not None:
                return (model_cls, arch)
715

716
717
718
        # 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):
719
720
721
722
723
724
725
            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)

726
        return self._raise_for_unsupported(architectures)
727

728
729
    def is_text_generation_model(
        self,
730
        architectures: Union[str, list[str]],
731
        model_config: ModelConfig,
732
    ) -> bool:
733
        model_cls, _ = self.inspect_model_cls(architectures, model_config)
734
        return model_cls.is_text_generation_model
735

736
    def is_pooling_model(
737
        self,
738
        architectures: Union[str, list[str]],
739
        model_config: ModelConfig,
740
    ) -> bool:
741
        model_cls, _ = self.inspect_model_cls(architectures, model_config)
742
        return model_cls.is_pooling_model
743

744
745
    def is_cross_encoder_model(
        self,
746
        architectures: Union[str, list[str]],
747
        model_config: ModelConfig,
748
    ) -> bool:
749
        model_cls, _ = self.inspect_model_cls(architectures, model_config)
750
        return model_cls.supports_cross_encoding
751

752
753
    def is_multimodal_model(
        self,
754
        architectures: Union[str, list[str]],
755
        model_config: ModelConfig,
756
    ) -> bool:
757
        model_cls, _ = self.inspect_model_cls(architectures, model_config)
758
        return model_cls.supports_multimodal
759

760
    def is_multimodal_raw_input_only_model(
761
762
        self,
        architectures: Union[str, list[str]],
763
        model_config: ModelConfig,
764
    ) -> bool:
765
        model_cls, _ = self.inspect_model_cls(architectures, model_config)
766
        return model_cls.supports_multimodal_raw_input_only
767

768
769
    def is_pp_supported_model(
        self,
770
        architectures: Union[str, list[str]],
771
        model_config: ModelConfig,
772
    ) -> bool:
773
        model_cls, _ = self.inspect_model_cls(architectures, model_config)
774
        return model_cls.supports_pp
775

776
777
    def model_has_inner_state(
        self,
778
        architectures: Union[str, list[str]],
779
        model_config: ModelConfig,
780
    ) -> bool:
781
        model_cls, _ = self.inspect_model_cls(architectures, model_config)
782
        return model_cls.has_inner_state
783

784
785
    def is_attention_free_model(
        self,
786
        architectures: Union[str, list[str]],
787
        model_config: ModelConfig,
788
    ) -> bool:
789
        model_cls, _ = self.inspect_model_cls(architectures, model_config)
790
        return model_cls.is_attention_free
791

792
793
    def is_hybrid_model(
        self,
794
        architectures: Union[str, list[str]],
795
        model_config: ModelConfig,
796
    ) -> bool:
797
        model_cls, _ = self.inspect_model_cls(architectures, model_config)
798
799
        return model_cls.is_hybrid

800
801
    def is_noops_model(
        self,
802
        architectures: Union[str, list[str]],
803
        model_config: ModelConfig,
804
    ) -> bool:
805
        model_cls, _ = self.inspect_model_cls(architectures, model_config)
806
807
        return model_cls.has_noops

808
809
    def is_transcription_model(
        self,
810
        architectures: Union[str, list[str]],
811
        model_config: ModelConfig,
812
    ) -> bool:
813
        model_cls, _ = self.inspect_model_cls(architectures, model_config)
814
815
        return model_cls.supports_transcription

816
817
818
    def is_transcription_only_model(
        self,
        architectures: Union[str, list[str]],
819
        model_config: ModelConfig,
820
    ) -> bool:
821
        model_cls, _ = self.inspect_model_cls(architectures, model_config)
822
823
        return model_cls.supports_transcription_only

824
825
    def is_v1_compatible(
        self,
826
        architectures: Union[str, list[str]],
827
        model_config: ModelConfig,
828
    ) -> bool:
829
        model_cls, _ = self.inspect_model_cls(architectures, model_config)
830
831
        return not model_cls.supports_v0_only

832
833

ModelRegistry = _ModelRegistry({
834
835
    model_arch:
    _LazyRegisteredModel(
836
837
838
839
840
841
842
843
844
845
        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:
846
847
848
849
850
    # 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")

851
        # `cloudpickle` allows pickling lambda functions directly
852
        import cloudpickle
853
        input_bytes = cloudpickle.dumps((fn, output_filepath))
854
855
856

        # cannot use `sys.executable __file__` here because the script
        # contains relative imports
857
858
859
        returned = subprocess.run(_SUBPROCESS_COMMAND,
                                  input=input_bytes,
                                  capture_output=True)
860
861
862
863
864
865
866
867
868

        # 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

869
        with open(output_filepath, "rb") as f:
870
871
872
873
874
875
876
877
878
879
880
            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()
881
882
883

    with open(output_file, "wb") as f:
        f.write(pickle.dumps(result))
884
885
886


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