91 lines
3.5 KiB
Python
91 lines
3.5 KiB
Python
import importlib
|
|
from typing import List, Optional, Type
|
|
|
|
import torch.nn as nn
|
|
|
|
from vllm.logger import init_logger
|
|
from vllm.utils import is_hip
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
# Architecture -> (module, class).
|
|
_MODELS = {
|
|
"AquilaModel": ("aquila", "AquilaForCausalLM"),
|
|
"AquilaForCausalLM": ("aquila", "AquilaForCausalLM"), # AquilaChat2
|
|
"BaiChuanForCausalLM": ("baichuan", "BaiChuanForCausalLM"), # baichuan-7b
|
|
"BaichuanForCausalLM": ("baichuan", "BaichuanForCausalLM"), # baichuan-13b
|
|
"BloomForCausalLM": ("bloom", "BloomForCausalLM"),
|
|
"ChatGLMModel": ("chatglm", "ChatGLMForCausalLM"),
|
|
"ChatGLMForConditionalGeneration": ("chatglm", "ChatGLMForCausalLM"),
|
|
"DeciLMForCausalLM": ("decilm", "DeciLMForCausalLM"),
|
|
"DeepseekForCausalLM": ("deepseek", "DeepseekForCausalLM"),
|
|
"FalconForCausalLM": ("falcon", "FalconForCausalLM"),
|
|
"GPT2LMHeadModel": ("gpt2", "GPT2LMHeadModel"),
|
|
"GPTBigCodeForCausalLM": ("gpt_bigcode", "GPTBigCodeForCausalLM"),
|
|
"GPTJForCausalLM": ("gpt_j", "GPTJForCausalLM"),
|
|
"GPTNeoXForCausalLM": ("gpt_neox", "GPTNeoXForCausalLM"),
|
|
"InternLMForCausalLM": ("internlm", "InternLMForCausalLM"),
|
|
"InternLM2ForCausalLM": ("internlm2", "InternLM2ForCausalLM"),
|
|
"LlamaForCausalLM": ("llama", "LlamaForCausalLM"),
|
|
# For decapoda-research/llama-*
|
|
"LLaMAForCausalLM": ("llama", "LlamaForCausalLM"),
|
|
"MistralForCausalLM": ("mistral", "MistralForCausalLM"),
|
|
"MixtralForCausalLM": ("mixtral", "MixtralForCausalLM"),
|
|
"QuantMixtralForCausalLM": ("mixtral_quant", "MixtralForCausalLM"),
|
|
# transformers's mpt class has lower case
|
|
"MptForCausalLM": ("mpt", "MPTForCausalLM"),
|
|
"MPTForCausalLM": ("mpt", "MPTForCausalLM"),
|
|
"OPTForCausalLM": ("opt", "OPTForCausalLM"),
|
|
"PhiForCausalLM": ("phi", "PhiForCausalLM"),
|
|
"QWenLMHeadModel": ("qwen", "QWenLMHeadModel"),
|
|
"Qwen2ForCausalLM": ("qwen2", "Qwen2ForCausalLM"),
|
|
"RWForCausalLM": ("falcon", "FalconForCausalLM"),
|
|
"StableLMEpochForCausalLM": ("stablelm", "StablelmForCausalLM"),
|
|
"YiForCausalLM": ("yi", "YiForCausalLM")
|
|
}
|
|
|
|
# Models not supported by ROCm.
|
|
_ROCM_UNSUPPORTED_MODELS = []
|
|
|
|
# Models partially supported by ROCm.
|
|
# Architecture -> Reason.
|
|
_ROCM_PARTIALLY_SUPPORTED_MODELS = {
|
|
"Qwen2ForCausalLM":
|
|
"Sliding window attention is not yet supported in ROCm's flash attention",
|
|
"MistralForCausalLM":
|
|
"Sliding window attention is not yet supported in ROCm's flash attention",
|
|
"MixtralForCausalLM":
|
|
"Sliding window attention is not yet supported in ROCm's flash attention",
|
|
}
|
|
|
|
|
|
class ModelRegistry:
|
|
|
|
@staticmethod
|
|
def load_model_cls(model_arch: str) -> Optional[Type[nn.Module]]:
|
|
if model_arch not in _MODELS:
|
|
return None
|
|
if is_hip():
|
|
if model_arch in _ROCM_UNSUPPORTED_MODELS:
|
|
raise ValueError(
|
|
f"Model architecture {model_arch} is not supported by "
|
|
"ROCm for now.")
|
|
if model_arch in _ROCM_PARTIALLY_SUPPORTED_MODELS:
|
|
logger.warning(
|
|
f"Model architecture {model_arch} is partially supported "
|
|
"by ROCm: " + _ROCM_PARTIALLY_SUPPORTED_MODELS[model_arch])
|
|
|
|
module_name, model_cls_name = _MODELS[model_arch]
|
|
module = importlib.import_module(
|
|
f"vllm.model_executor.models.{module_name}")
|
|
return getattr(module, model_cls_name, None)
|
|
|
|
@staticmethod
|
|
def get_supported_archs() -> List[str]:
|
|
return list(_MODELS.keys())
|
|
|
|
|
|
__all__ = [
|
|
"ModelRegistry",
|
|
]
|