vLLM/tests/models/registry.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

from collections.abc import Mapping, Set
from dataclasses import dataclass, field
from typing import Any, Literal

import pytest
from packaging.version import Version
from transformers import PreTrainedConfig
from transformers import __version__ as TRANSFORMERS_VERSION

from vllm.config.model import ModelDType, TokenizerMode


@dataclass(frozen=True)
class _HfExamplesInfo:
    default: str
    """The default model to use for testing this architecture."""

    extras: Mapping[str, str] = field(default_factory=dict)
    """Extra models to use for testing this architecture."""

    tokenizer: str | None = None
    """Set the tokenizer to load for this architecture."""

    tokenizer_mode: TokenizerMode | str = "auto"
    """Set the tokenizer type for this architecture."""

    speculative_model: str | None = None
    """
    The default model to use for testing this architecture, which is only used
    for speculative decoding.
    """

    speculative_method: str | None = None
    """
    The method to use for speculative decoding.
    """

    min_transformers_version: str | None = None
    """
    The minimum version of HF Transformers that is required to run this model.
    """

    max_transformers_version: str | None = None
    """
    The maximum version of HF Transformers that this model runs on.
    """

    transformers_version_reason: dict[Literal["vllm", "hf"], str] | None = None
    """
    The type and reason to skip test for the minimum/maximum version requirement.
    vllm: skip all vLLM tests if the version requirement is not met.
    hf: only skip tests that uses HF runner if the version requirement is not met.
    """

    require_embed_inputs: bool = False
    """
    If `True`, enables prompt and multi-modal embedding inputs while
    disabling tokenization.
    """

    dtype: ModelDType = "auto"
    """
    The data type for the model weights and activations.
    """

    enforce_eager: bool = False
    """
    Whether to enforce eager execution. If True, we will
    disable CUDA graph and always execute the model in eager mode.
    If False, we will use CUDA graph and eager execution in hybrid.
    """

    enable_prefix_caching: bool = True
    """
    Whether to enable prefix caching for the model. If True, we will test the model with
    prefix caching enabled. If False, we will test the model without prefix caching.
    """

    is_available_online: bool = True
    """
    Set this to `False` if the name of this architecture no longer exists on
    the HF repo. To maintain backwards compatibility, we have not removed them
    from the main model registry, so without this flag the registry tests will
    fail.
    """

    trust_remote_code: bool = False
    """The `trust_remote_code` level required to load the model."""

    hf_overrides: dict[str, Any] = field(default_factory=dict)
    """The `hf_overrides` required to load the model."""

    max_model_len: int | None = None
    """
    The maximum model length to use for this model. Some models default to a
    length that is too large to fit into memory in CI.
    """

    max_num_batched_tokens: int | None = None
    """
    The maximum number of tokens to be processed in a single batch.
    """

    revision: str | None = None
    """
    The specific revision (commit hash, tag, or branch) to use for the model.
    If not specified, the default revision will be used.
    """

    max_num_seqs: int | None = None
    """Maximum number of sequences to be processed in a single iteration."""

    use_original_num_layers: bool = False
    """
    If True, use the original number of layers from the model config
    instead of minimal layers for testing.
    """

    def check_transformers_version(
        self,
        *,
        on_fail: Literal["error", "skip", "return"],
        check_version_reason: Literal["vllm", "hf"] = "hf",
        check_min_version: bool = True,
        check_max_version: bool = True,
    ) -> str | None:
        """If the installed transformers version does not meet the requirements,
        perform the given action.
        """
        if (
            self.min_transformers_version is None
            and self.max_transformers_version is None
        ):
            return None

        current_version = TRANSFORMERS_VERSION
        cur_base_version = Version(current_version).base_version
        min_version = self.min_transformers_version
        max_version = self.max_transformers_version
        msg = f"`transformers=={current_version}` installed, but `transformers"
        # Only check the base version for the min/max version, otherwise preview
        # models cannot be run because `x.yy.0.dev0`<`x.yy.0`
        if min_version and Version(cur_base_version) < Version(min_version):
            is_version_valid = False
            should_check_version = check_min_version
            msg += f">={min_version}` is required to run this model."
        elif max_version and Version(cur_base_version) > Version(max_version):
            is_version_valid = False
            should_check_version = check_max_version
            msg += f"<={max_version}` is required to run this model."
        else:
            is_version_valid = True
            should_check_version = False

        # Reasons explain a known incompatibility with a violated version
        # bound. They should not skip models when the installed version is
        # already within the requested range.
        is_reason_applicable = (
            not is_version_valid
            and self.transformers_version_reason is not None
            and check_version_reason in self.transformers_version_reason
        )
        is_transformers_valid = is_version_valid or (
            not should_check_version and not is_reason_applicable
        )
        if is_transformers_valid:
            return None
        elif self.transformers_version_reason:
            for reason_type, reason in self.transformers_version_reason.items():
                msg += f" Reason({reason_type}): {reason}"

        if on_fail == "error":
            raise RuntimeError(msg)
        elif on_fail == "skip":
            pytest.skip(msg)

        return msg

    def check_available_online(
        self,
        *,
        on_fail: Literal["error", "skip"],
    ) -> None:
        """If the model is not available online, perform the given action."""
        if not self.is_available_online:
            msg = "Model is not available online"

            if on_fail == "error":
                raise RuntimeError(msg)
            else:
                pytest.skip(msg)


_TEXT_GENERATION_EXAMPLE_MODELS = {
    "ArceeForCausalLM": _HfExamplesInfo("arcee-ai/AFM-4.5B-Base"),
    "CwmForCausalLM": _HfExamplesInfo("facebook/cwm"),
    "FlexOlmoForCausalLM": _HfExamplesInfo("allenai/Flex-reddit-2x7B-1T"),
    "GlmForCausalLM": _HfExamplesInfo("zai-org/glm-4-9b-chat-hf"),
    "GPTBigCodeForCausalLM": _HfExamplesInfo(
        "bigcode/starcoder",
        extras={
            "tiny": "bigcode/tiny_starcoder_py",
            "santacoder": "bigcode/gpt_bigcode-santacoder",
        },
    ),
    "GPTNeoXForCausalLM": _HfExamplesInfo(
        "EleutherAI/pythia-70m", {"1b": "EleutherAI/pythia-1.4b"}
    ),
    "HunYuanDenseV1ForCausalLM": _HfExamplesInfo("tencent/Hunyuan-7B-Instruct"),
    "HunYuanMoEV1ForCausalLM": _HfExamplesInfo("tencent/Hunyuan-A13B-Instruct"),
    "Jais2ForCausalLM": _HfExamplesInfo("inceptionai/Jais-2-8B-Chat"),
    "Kolibri1ForCausalLM": _HfExamplesInfo(
        "Aleph-Alpha/Kolibri-1", min_transformers_version="5.20.0"
    ),
    "NanbeigeForCausalLM": _HfExamplesInfo(
        "Nanbeige/Nanbeige4.2-3B",
        trust_remote_code=True,
    ),
    "OlmoForCausalLM": _HfExamplesInfo("allenai/OLMo-1B-hf"),
    "Olmo2ForCausalLM": _HfExamplesInfo("allenai/OLMo-2-0425-1B"),
    "Olmo3ForCausalLM": _HfExamplesInfo("allenai/Olmo-3-7B-Instruct"),
    "PhiForCausalLM": _HfExamplesInfo("microsoft/phi-2"),
    "Qwen2ForCausalLM": _HfExamplesInfo(
        "Qwen/Qwen2-0.5B-Instruct",
        extras={
            "2.5": "Qwen/Qwen2.5-0.5B-Instruct",
            "2.5-1.5B": "Qwen/Qwen2.5-1.5B-Instruct",
        },
    ),
    "Qwen2MoeForCausalLM": _HfExamplesInfo("Qwen/Qwen1.5-MoE-A2.7B-Chat"),
    "Qwen3ForCausalLM": _HfExamplesInfo("Qwen/Qwen3-8B"),
    "Qwen3_5ForCausalLM": _HfExamplesInfo("codecho/Qwen3.5-0.8B-text-only"),
    "Qwen3_5MoeForCausalLM": _HfExamplesInfo(
        "codecho/Qwen3.5-35B-A3B-text-only",
        extras={"native-prefix": "imdatta0/small_qwen3_5_20b"},
        max_model_len=4096,
    ),
    "Qwen4ExpForCausalLM": _HfExamplesInfo("", is_available_online=False),
    "MellumForCausalLM": _HfExamplesInfo("JetBrains/Mellum2-12B-A2.5B-Base"),
    "Qwen3NextForCausalLM": _HfExamplesInfo(
        "Qwen/Qwen3-Next-80B-A3B-Instruct",
        extras={"tiny-random": "tiny-random/qwen3-next-moe"},
    ),
    "SeedOssForCausalLM": _HfExamplesInfo(
        "ByteDance-Seed/Seed-OSS-36B-Instruct",
        trust_remote_code=True,
    ),
    "SmolLM3ForCausalLM": _HfExamplesInfo("HuggingFaceTB/SmolLM3-3B"),
    "Starcoder2ForCausalLM": _HfExamplesInfo("bigcode/starcoder2-3b"),
    "VaultGemmaForCausalLM": _HfExamplesInfo("google/vaultgemma-1b"),
}

_EMBEDDING_EXAMPLE_MODELS = {
    "Qwen2Model": _HfExamplesInfo("ssmits/Qwen2-7B-Instruct-embed-base"),
    "Qwen2VLForConditionalGeneration": _HfExamplesInfo("MrLight/dse-qwen2-2b-mrl-v1"),
}

_LATE_INTERACTION_EXAMPLE_MODELS = {
}


_REWARD_EXAMPLE_MODELS = {
}

_TOKEN_CLASSIFICATION_EXAMPLE_MODELS = {
}

_SEQUENCE_CLASSIFICATION_EXAMPLE_MODELS = {
}

_AUTOMATIC_CONVERTED_MODELS = {
}

_MULTIMODAL_EXAMPLE_MODELS = {
    "Emu3ForConditionalGeneration": _HfExamplesInfo("BAAI/Emu3-Chat-hf"),
    "HunYuanVLForConditionalGeneration": _HfExamplesInfo(
        "tencent/HunyuanOCR",
    ),
    "Qwen2VLForConditionalGeneration": _HfExamplesInfo("Qwen/Qwen2-VL-2B-Instruct"),
    "Qwen3VLForConditionalGeneration": _HfExamplesInfo(
        "Qwen/Qwen3-VL-4B-Instruct",
        max_model_len=4096,
    ),
    "Qwen3_5ForConditionalGeneration": _HfExamplesInfo(
        "Qwen/Qwen3.5-0.8B",
        extras={"4b": "Qwen/Qwen3.5-4B"},
        max_model_len=4096,
    ),
    "Qwen3_5MoeForConditionalGeneration": _HfExamplesInfo(
        "Qwen/Qwen3.5-35B-A3B",
        max_model_len=4096,
    ),
    "Qwen4ExpForConditionalGeneration": _HfExamplesInfo(
        "",
        is_available_online=False,
    ),
    "VibeVoiceAsrForConditionalGeneration": _HfExamplesInfo(
        "microsoft/VibeVoice-ASR-HF",
    ),
}


_SPECULATIVE_DECODING_EXAMPLE_MODELS = {
    # [DFlash]
    "DFlashDraftModel": _HfExamplesInfo(
        "Qwen/Qwen3-4B",
        speculative_model="z-lab/Qwen3-4B-DFlash-b16",
        use_original_num_layers=True,  # Need all layers since DFlash has >1 layer,
        max_model_len=8192,  # Reduce max len to ensure test runs in low-VRAM CI env
        max_num_seqs=32,
    ),
    "MuseGlimmerAssistantModel": _HfExamplesInfo(
        "meta-models/Muse-Glimmer-30B",
        speculative_model="meta-models/Muse-Glimmer-30B-assistant",
        max_model_len=8192,  # Reduce max len to ensure test runs in low-VRAM CI env
        max_num_seqs=32,
    ),
    "DFlashMuseGlimmerAssistantModel": _HfExamplesInfo(
        "meta-models/Muse-Glimmer-30B",
        speculative_model="meta-models/Muse-Glimmer-30B-assistant",
        max_model_len=8192,  # Reduce max len to ensure test runs in low-VRAM CI env
        max_num_seqs=32,
    ),
    "Qwen3_5MTP": _HfExamplesInfo(
        "Qwen/Qwen3.5-0.8B",
        speculative_model="Qwen/Qwen3.5-0.8B",
    ),
    "Qwen3_5MoeMTP": _HfExamplesInfo(
        "Qwen/Qwen3.5-35B-A3B",
        speculative_model="Qwen/Qwen3.5-35B-A3B",
    ),
    "Qwen4ExpMTP": _HfExamplesInfo("", is_available_online=False),
}

_TRANSFORMERS_BACKEND_MODELS = {
    "TransformersEmbeddingModel": _HfExamplesInfo("BAAI/bge-base-en-v1.5"),
    "TransformersForSequenceClassification": _HfExamplesInfo(
        "papluca/xlm-roberta-base-language-detection",
    ),
    "TransformersForCausalLM": _HfExamplesInfo(
        "hmellor/Ilama-3.2-1B", trust_remote_code=True
    ),
    "TransformersMultiModalForCausalLM": _HfExamplesInfo("BAAI/Emu3-Chat-hf"),
    "TransformersMoEForCausalLM": _HfExamplesInfo("allenai/OLMoE-1B-7B-0924"),
    "TransformersMultiModalMoEForCausalLM": _HfExamplesInfo(
        "Qwen/Qwen3-VL-30B-A3B-Instruct"
    ),
    "TransformersMoEEmbeddingModel": _HfExamplesInfo("Qwen/Qwen3-30B-A3B"),
    "TransformersMoEForSequenceClassification": _HfExamplesInfo("Qwen/Qwen3-30B-A3B"),
    "TransformersMultiModalEmbeddingModel": _HfExamplesInfo("google/gemma-3-4b-it"),
    "TransformersMultiModalForSequenceClassification": _HfExamplesInfo(
        "google/gemma-3-4b-it"
    ),
}

_EXAMPLE_MODELS = {
    **_TEXT_GENERATION_EXAMPLE_MODELS,
    **_EMBEDDING_EXAMPLE_MODELS,
    **_LATE_INTERACTION_EXAMPLE_MODELS,
    **_REWARD_EXAMPLE_MODELS,
    **_TOKEN_CLASSIFICATION_EXAMPLE_MODELS,
    **_SEQUENCE_CLASSIFICATION_EXAMPLE_MODELS,
    **_MULTIMODAL_EXAMPLE_MODELS,
    **_SPECULATIVE_DECODING_EXAMPLE_MODELS,
    **_TRANSFORMERS_BACKEND_MODELS,
}


class HfExampleModels:
    def __init__(self, hf_models: Mapping[str, _HfExamplesInfo]) -> None:
        super().__init__()

        self.hf_models = hf_models

    def get_supported_archs(self) -> Set[str]:
        return self.hf_models.keys()

    def get_hf_info(self, model_arch: str) -> _HfExamplesInfo:
        try:
            return self.hf_models[model_arch]
        except KeyError:
            raise ValueError(
                f"No example model defined for {model_arch}; please update this file."
            ) from None

    def find_hf_info(self, model_id: str) -> _HfExamplesInfo:
        for info in self.hf_models.values():
            if info.default == model_id:
                return info

        # Fallback to extras
        for info in self.hf_models.values():
            if any(extra == model_id for extra in info.extras.values()):
                return info

        raise ValueError(
            f"No example model defined for {model_id}; please update this file."
        )


HF_EXAMPLE_MODELS = HfExampleModels(_EXAMPLE_MODELS)
AUTO_EXAMPLE_MODELS = HfExampleModels(_AUTOMATIC_CONVERTED_MODELS)