test_common.py 5.92 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
3
4
import os
from typing import Optional

Woosuk Kwon's avatar
Woosuk Kwon committed
5
import pytest
6
7
8
import torch

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

10
from ....utils import large_gpu_mark
11
from ...registry import HF_EXAMPLE_MODELS
12
from ...utils import check_logprobs_close
13

14
15
16
17
18
19
# 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"]

20
21
22
23
24
25
26
27
# 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",
28
    "Qwen/Qwen-7B-Chat",
29
    "Qwen/Qwen2.5-0.5B-Instruct",
30
    "TitanML/tiny-mixtral",
31
32
]

Woosuk Kwon's avatar
Woosuk Kwon committed
33

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

96
97
98
    model_info = HF_EXAMPLE_MODELS.find_hf_info(model)
    model_info.check_available_online(on_fail="skip")
    model_info.check_transformers_version(on_fail="skip")
99

100
101
    if model in REQUIRES_V0:
        monkeypatch.setenv("VLLM_USE_V1", "0")
102

103
104
105
106
107
108
109
110
111
    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.")

112
113
    use_prompt_embeds = os.getenv("VLLM_USE_V1") == "0"

114
    with hf_runner(model) as hf_model:
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
119
120
        prompt_embeds: Optional[list[torch.Tensor]] = ([] if use_prompt_embeds
                                                       else None)

121
122
123
124
125
126
127
128
129
130
        prompt_token_ids = []
        for prompt in example_prompts:
            token_ids = hf_model.tokenizer(prompt,
                                           return_tensors="pt").input_ids.to(
                                               hf_model.model.device)
            prompt_token_ids.append(token_ids)
            if prompt_embeds is not None:
                prompt_embeds.append(hf_model.model.get_input_embeddings()(
                    token_ids).squeeze(0))

131
132
133
134
135
136
    with vllm_runner(
            model,
            tokenizer_name=model_info.tokenizer or model,
            tokenizer_mode=model_info.tokenizer_mode,
            trust_remote_code=model_info.trust_remote_code,
            max_num_seqs=2,
137
            enable_prompt_embeds=use_prompt_embeds,
138
    ) as vllm_model:
139
140
        vllm_outputs = vllm_model.generate_greedy_logprobs(
            example_prompts, max_tokens, num_logprobs)
141
142
143
        if prompt_embeds is not None:
            vllm_outputs_from_embeds = vllm_model.generate_greedy_logprobs(
                prompt_embeds, max_tokens, num_logprobs)
144

145
    check_logprobs_close(
146
147
148
149
150
        outputs_0_lst=hf_outputs,
        outputs_1_lst=vllm_outputs,
        name_0="hf",
        name_1="vllm",
    )
151
152
153
154
155
156
157
158
    if prompt_embeds is not None:
        check_logprobs_close(
            outputs_0_lst=vllm_outputs,
            outputs_1_lst=vllm_outputs_from_embeds,
            name_0="vllm",
            name_1="vllm_from_embeds",
        )

159
160
161
162
163
164
165
    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()