test_models.py 5.22 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
Woosuk Kwon's avatar
Woosuk Kwon committed
2
3
"""Compare the outputs of HF and vLLM when using greedy sampling.

4
Run `pytest tests/models/test_models.py`.
Woosuk Kwon's avatar
Woosuk Kwon committed
5
"""
6

Woosuk Kwon's avatar
Woosuk Kwon committed
7
import pytest
8
import os
9
10
11
import torch

from vllm.platforms import current_platform
Woosuk Kwon's avatar
Woosuk Kwon committed
12

13
from ...utils import check_logprobs_close
zhuwenwen's avatar
zhuwenwen committed
14
from ....utils import models_path_prefix
Woosuk Kwon's avatar
Woosuk Kwon committed
15

16
17
18
19
20
21
# These have unsupported head_dim for FA. We do not
# not have a clean way to fall back, so we fail with
# a clear msg when it happens.
# https://github.com/vllm-project/vllm/issues/14524
REQUIRES_V0 = ["microsoft/phi-2", "stabilityai/stablelm-3b-4e1t"]

22
23
24
25
26
27
28
29
30
31
32
33
# This list contains the model that are using AITER kernel.
# Skip model that are not using AITER tests.
# When more AITER kernels are added, this list will not be
# needed as all the models will be calling AITER kernels
# in parts of the operators
AITER_MODEL_LIST = [
    "meta-llama/Llama-3.2-1B-Instruct",
    "openbmb/MiniCPM3-4B",
    "Qwen/Qwen-7B",
    "Qwen/Qwen2.5-0.5B-Instruct",
    "ehristoforu/Falcon3-MoE-2x7B-Insruct",
]
Woosuk Kwon's avatar
Woosuk Kwon committed
34
35


36
# @maybe_test_rocm_aiter
37
38
39
40
@pytest.mark.parametrize(
    "model",
    [
        pytest.param(
zhuwenwen's avatar
zhuwenwen committed
41
            os.path.join(models_path_prefix, "bigscience/bloom-560m"),  # bloom - testing alibi slopes
42
43
44
            marks=[pytest.mark.core_model, pytest.mark.cpu_model],
        ),
        pytest.param(
zhuwenwen's avatar
zhuwenwen committed
45
            os.path.join(models_path_prefix, "openai-community/gpt2"),  # gpt2
46
47
            marks=[pytest.mark.core_model, pytest.mark.cpu_model],
        ),
zhuwenwen's avatar
zhuwenwen committed
48
49
50
        pytest.param(os.path.join(models_path_prefix, "Milos/slovak-gpt-j-405M")),  # gptj
        pytest.param(os.path.join(models_path_prefix, "bigcode/tiny_starcoder_py")),  # gpt_bigcode
        pytest.param(os.path.join(models_path_prefix, "EleutherAI/pythia-70m")),  # gpt_neox
51
        pytest.param(
zhuwenwen's avatar
zhuwenwen committed
52
            os.path.join(models_path_prefix, "google/gemma-1.1-2b-it"),  # gemma
53
54
            marks=[pytest.mark.core_model, pytest.mark.cpu_model],
        ),
55
        pytest.param(
zhuwenwen's avatar
zhuwenwen committed
56
            os.path.join(models_path_prefix, "THUDM/chatglm3-6b"),  # chatglm (text-only)
57
        ),
58
        pytest.param(
zhuwenwen's avatar
zhuwenwen committed
59
            os.path.join(models_path_prefix, "meta-llama/Llama-3.2-1B-Instruct"),  # llama
60
61
62
            marks=[pytest.mark.core_model, pytest.mark.cpu_model],
        ),
        pytest.param(
zhuwenwen's avatar
zhuwenwen committed
63
            os.path.join(models_path_prefix, "openbmb/MiniCPM3-4B"),
64
65
66
67
            # fused_moe not supported on CPU
            marks=[pytest.mark.core_model],
        ),
        pytest.param(
zhuwenwen's avatar
zhuwenwen committed
68
            os.path.join(models_path_prefix, "facebook/opt-125m"),  # opt
69
70
71
            marks=[pytest.mark.core_model, pytest.mark.cpu_model],
        ),
        pytest.param(
zhuwenwen's avatar
zhuwenwen committed
72
            os.path.join(models_path_prefix, "microsoft/phi-2"),  # phi
73
74
            marks=[pytest.mark.core_model],
        ),
75
        pytest.param(
zhuwenwen's avatar
zhuwenwen committed
76
            os.path.join(models_path_prefix, "Qwen/Qwen-7B"),  # qwen (text-only)
77
        ),
78
        pytest.param(
zhuwenwen's avatar
zhuwenwen committed
79
            os.path.join(models_path_prefix, "Qwen/Qwen2.5-0.5B-Instruct"),  # qwen2
80
81
            marks=[pytest.mark.core_model],
        ),
zhuwenwen's avatar
zhuwenwen committed
82
83
        pytest.param(os.path.join(models_path_prefix, "stabilityai/stablelm-3b-4e1t")),  # stablelm
        pytest.param(os.path.join(models_path_prefix, "bigcode/starcoder2-3b")),  # starcoder2
84
        pytest.param(
zhuwenwen's avatar
zhuwenwen committed
85
            os.path.join(models_path_prefix, "ehristoforu/Falcon3-MoE-2x7B-Insruct"),  # mixtral
86
87
            marks=[pytest.mark.cpu_model],
        )
88
89
    ])
@pytest.mark.parametrize("dtype", ["half"])
90
91
@pytest.mark.parametrize("max_tokens", [32])
@pytest.mark.parametrize("num_logprobs", [5])
92
93
94
95
96
97
@pytest.mark.parametrize(
    "use_rocm_aiter", [True, False] if current_platform.is_rocm() else [False])
def test_models(hf_runner, vllm_runner, example_prompts, model: str,
                dtype: str, max_tokens: int, num_logprobs: int,
                use_rocm_aiter: bool, monkeypatch) -> None:

98
99
    if model in REQUIRES_V0:
        monkeypatch.setenv("VLLM_USE_V1", "0")
100

101
102
103
104
105
106
107
108
    if use_rocm_aiter and (model in AITER_MODEL_LIST):
        monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1")
    elif use_rocm_aiter and model not in AITER_MODEL_LIST:
        # Skip model that are not using AITER tests.
        # When more AITER kernels are added, this list will not be
        # needed as all the models will be calling AITER kernels
        # in parts of the operators
        pytest.skip(f"Skipping '{model}' model test with AITER kernel.")
109

110
    with hf_runner(model, dtype=dtype) as hf_model:
111
112
113
114
        if model.startswith("THUDM/chatglm3"):
            hf_model.model.get_output_embeddings = lambda: \
                hf_model.model.transformer.output_layer

115
116
        hf_outputs = hf_model.generate_greedy_logprobs_limit(
            example_prompts, max_tokens, num_logprobs)
Woosuk Kwon's avatar
Woosuk Kwon committed
117

118
    with vllm_runner(model, dtype=dtype) as vllm_model:
119
120
        vllm_outputs = vllm_model.generate_greedy_logprobs(
            example_prompts, max_tokens, num_logprobs)
121

122
    check_logprobs_close(
123
124
125
126
127
        outputs_0_lst=hf_outputs,
        outputs_1_lst=vllm_outputs,
        name_0="hf",
        name_1="vllm",
    )
128
129
130
131
132
133
134
    if use_rocm_aiter:
        # this is to ensure that vllm engine
        # has deallocated the memory before running the next
        # unit tests. On ROCm, when using AITER
        # the memory might not be deallocated completely
        # before running the next test case
        torch.cuda.synchronize()