conftest.py 8.74 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3

4
5
import tempfile
from collections import OrderedDict
6
from unittest.mock import MagicMock, patch
7
8
9
10
11
12
13
14

import pytest
import torch
import torch.nn as nn
from huggingface_hub import snapshot_download

import vllm
from vllm.config import LoRAConfig
15
from vllm.distributed import (cleanup_dist_env_and_memory,
16
17
                              init_distributed_environment,
                              initialize_model_parallel)
18
19
20
from vllm.model_executor.layers.linear import (ColumnParallelLinear,
                                               MergedColumnParallelLinear,
                                               RowParallelLinear)
21
22
from vllm.model_executor.layers.logits_processor import LogitsProcessor
from vllm.model_executor.layers.sampler import Sampler
23
from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
24
from vllm.model_executor.model_loader import get_model
25
from vllm.model_executor.models.interfaces import SupportsLoRA
26
from vllm.platforms import current_platform
27

28

29
30
31
32
33
34
35
@pytest.fixture()
def should_do_global_cleanup_after_test(request) -> bool:
    """Allow subdirectories to skip global cleanup by overriding this fixture.
    This can provide a ~10x speedup for non-GPU unit tests since they don't need
    to initialize torch.
    """

36
    return not request.node.get_closest_marker("skip_global_cleanup")
37
38


39
@pytest.fixture(autouse=True)
40
def cleanup_fixture(should_do_global_cleanup_after_test: bool):
41
    yield
42
    if should_do_global_cleanup_after_test:
43
        cleanup_dist_env_and_memory(shutdown_ray=True)
44
45
46
47


@pytest.fixture
def dist_init():
48
    temp_file = tempfile.mkstemp()[1]
49
50

    backend = "nccl"
51
    if current_platform.is_cpu() or current_platform.is_tpu():
52
53
54
55
56
57
58
        backend = "gloo"

    init_distributed_environment(world_size=1,
                                 rank=0,
                                 distributed_init_method=f"file://{temp_file}",
                                 local_rank=0,
                                 backend=backend)
59
60
    initialize_model_parallel(1, 1)
    yield
61
    cleanup_dist_env_and_memory(shutdown_ray=True)
62
63
64
65
66
67


@pytest.fixture
def dist_init_torch_only():
    if torch.distributed.is_initialized():
        return
68
69
70
71
    backend = "nccl"
    if current_platform.is_cpu():
        backend = "gloo"

72
    temp_file = tempfile.mkstemp()[1]
73
74
75
76
    torch.distributed.init_process_group(world_size=1,
                                         rank=0,
                                         init_method=f"file://{temp_file}",
                                         backend=backend)
77
78


79
80
81
82
class DummyLoRAModel(nn.Sequential, SupportsLoRA):
    pass


83
84
@pytest.fixture
def dummy_model() -> nn.Module:
85
    model = DummyLoRAModel(
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
        OrderedDict([
            ("dense1", ColumnParallelLinear(764, 100)),
            ("dense2", RowParallelLinear(100, 50)),
            (
                "layer1",
                nn.Sequential(
                    OrderedDict([
                        ("dense1", ColumnParallelLinear(100, 10)),
                        ("dense2", RowParallelLinear(10, 50)),
                    ])),
            ),
            ("act2", nn.ReLU()),
            ("output", ColumnParallelLinear(50, 10)),
            ("outact", nn.Sigmoid()),
            # Special handling for lm_head & sampler
            ("lm_head", ParallelLMHead(512, 10)),
102
103
            ("logits_processor", LogitsProcessor(512)),
            ("sampler", Sampler())
104
105
        ]))
    model.config = MagicMock()
106
    model.embedding_modules = {"lm_head": "lm_head"}
107
108
109
110
111
    return model


@pytest.fixture
def dummy_model_gate_up() -> nn.Module:
112
    model = DummyLoRAModel(
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
        OrderedDict([
            ("dense1", ColumnParallelLinear(764, 100)),
            ("dense2", RowParallelLinear(100, 50)),
            (
                "layer1",
                nn.Sequential(
                    OrderedDict([
                        ("dense1", ColumnParallelLinear(100, 10)),
                        ("dense2", RowParallelLinear(10, 50)),
                    ])),
            ),
            ("act2", nn.ReLU()),
            ("gate_up_proj", MergedColumnParallelLinear(50, [5, 5])),
            ("outact", nn.Sigmoid()),
            # Special handling for lm_head & sampler
            ("lm_head", ParallelLMHead(512, 10)),
129
130
            ("logits_processor", LogitsProcessor(512)),
            ("sampler", Sampler())
131
132
        ]))
    model.config = MagicMock()
133
134
135
136
137
138
139
    model.packed_modules_mapping = {
        "gate_up_proj": [
            "gate_proj",
            "up_proj",
        ],
    }
    model.embedding_modules = {"lm_head": "lm_head"}
140
141
142
    return model


143
144
145
146
147
148
@pytest.fixture(scope="session")
def llama_2_7b_base_huggingface_id():
    # used as a base model for testing with sql lora adapter
    return "meta-llama/Llama-2-7b-hf"


149
@pytest.fixture(scope="session")
150
151
152
153
154
155
156
157
def sql_lora_huggingface_id():
    # huggingface repo id is used to test lora runtime downloading.
    return "yard1/llama-2-7b-sql-lora-test"


@pytest.fixture(scope="session")
def sql_lora_files(sql_lora_huggingface_id):
    return snapshot_download(repo_id=sql_lora_huggingface_id)
158
159


Terry's avatar
Terry committed
160
161
@pytest.fixture(scope="session")
def mixtral_lora_files():
162
163
164
    # Note: this module has incorrect adapter_config.json to test
    # https://github.com/vllm-project/vllm/pull/5909/files.
    return snapshot_download(repo_id="SangBinCho/mixtral-lora")
Terry's avatar
Terry committed
165
166


167
168
169
170
171
@pytest.fixture(scope="session")
def gemma_lora_files():
    return snapshot_download(repo_id="wskwon/gemma-7b-test-lora")


172
173
174
175
176
177
178
179
180
181
@pytest.fixture(scope="session")
def chatglm3_lora_files():
    return snapshot_download(repo_id="jeeejeee/chatglm3-text2sql-spider")


@pytest.fixture(scope="session")
def baichuan_lora_files():
    return snapshot_download(repo_id="jeeejeee/baichuan7b-text2sql-spider")


182
183
184
185
186
187
@pytest.fixture(scope="session")
def baichuan_zero_lora_files():
    # all the lora_B weights are initialized to zero.
    return snapshot_download(repo_id="jeeejeee/baichuan7b-zero-init")


188
189
190
191
192
@pytest.fixture(scope="session")
def baichuan_regex_lora_files():
    return snapshot_download(repo_id="jeeejeee/baichuan-7b-lora-zero-regex")


193
194
195
196
197
@pytest.fixture(scope="session")
def ilama_lora_files():
    return snapshot_download(repo_id="jeeejeee/ilama-text2sql-spider")


198
199
200
201
202
@pytest.fixture(scope="session")
def minicpmv_lora_files():
    return snapshot_download(repo_id="jeeejeee/minicpmv25-lora-pokemon")


203
204
205
206
207
@pytest.fixture(scope="session")
def qwen2vl_lora_files():
    return snapshot_download(repo_id="jeeejeee/qwen2-vl-lora-pokemon")


208
209
210
211
212
213
@pytest.fixture(scope="session")
def qwen25vl_base_huggingface_id():
    # used as a base model for testing with qwen25vl lora adapter
    return "Qwen/Qwen2.5-VL-3B-Instruct"


214
215
216
217
218
@pytest.fixture(scope="session")
def qwen25vl_lora_files():
    return snapshot_download(repo_id="jeeejeee/qwen25-vl-lora-pokemon")


219
220
221
222
223
@pytest.fixture(scope="session")
def tinyllama_lora_files():
    return snapshot_download(repo_id="jashing/tinyllama-colorist-lora")


224
225
226
227
228
@pytest.fixture(scope="session")
def phi2_lora_files():
    return snapshot_download(repo_id="isotr0py/phi-2-test-sql-lora")


229
230
231
232
233
@pytest.fixture(scope="session")
def long_context_lora_files_16k_1():
    return snapshot_download(repo_id="SangBinCho/long_context_16k_testing_1")


234
@pytest.fixture
235
def llama_2_7b_engine_extra_embeddings():
236
    cleanup_dist_env_and_memory(shutdown_ray=True)
237
238
    get_model_old = get_model

239
240
241
242
    def get_model_patched(**kwargs):
        kwargs["vllm_config"].lora_config = LoRAConfig(max_loras=4,
                                                       max_lora_rank=8)
        return get_model_old(**kwargs)
243
244
245
246
247

    with patch("vllm.worker.model_runner.get_model", get_model_patched):
        engine = vllm.LLM("meta-llama/Llama-2-7b-hf", enable_lora=False)
    yield engine.llm_engine
    del engine
248
    cleanup_dist_env_and_memory(shutdown_ray=True)
249
250
251


@pytest.fixture
252
def llama_2_7b_model_extra_embeddings(llama_2_7b_engine_extra_embeddings):
253
254
    yield (llama_2_7b_engine_extra_embeddings.model_executor.driver_worker.
           model_runner.model)
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271


@pytest.fixture(params=[True, False])
def run_with_both_engines_lora(request, monkeypatch):
    # Automatically runs tests twice, once with V1 and once without
    use_v1 = request.param
    # Tests decorated with `@skip_v1` are only run without v1
    skip_v1 = request.node.get_closest_marker("skip_v1")

    if use_v1:
        if skip_v1:
            pytest.skip("Skipping test on vllm V1")
        monkeypatch.setenv('VLLM_USE_V1', '1')
    else:
        monkeypatch.setenv('VLLM_USE_V1', '0')

    yield
272
273
274
275
276


@pytest.fixture
def reset_default_device():
    """
277
278
    Some tests, such as `test_punica_ops.py`, explicitly set the
    default device, which can affect subsequent tests. Adding this fixture
279
280
281
282
283
    helps avoid this problem.
    """
    original_device = torch.get_default_device()
    yield
    torch.set_default_device(original_device)