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

Mor Zusman's avatar
Mor Zusman committed
4
5
import pytest

6
from tests.utils import multi_gpu_test
7
from vllm.engine.arg_utils import EngineArgs
8
from vllm.sampling_params import SamplingParams
9

10
11
from ...utils import check_logprobs_close, check_outputs_equal

12
13
14
# Mark all tests as hybrid
pytestmark = pytest.mark.hybrid_model

15
16
17
18
19
20
21
22
# NOTE: The first model in each list is taken as the primary model,
# meaning that it will be used in all tests in this file
# The rest of the models will only be tested by test_models

SSM_MODELS = [
    "state-spaces/mamba-130m-hf",
    "tiiuae/falcon-mamba-tiny-dev",
    # TODO: Compare to a Mamba2 model. The HF transformers implementation of
Chen Zhang's avatar
Chen Zhang committed
23
24
    # Mamba2 is buggy for Codestral as it doesn't handle n_groups, so the test
    # doesn't compare vLLM output with HF output.
25
    # See https://github.com/huggingface/transformers/pull/35943
Chen Zhang's avatar
Chen Zhang committed
26
    "mistralai/Mamba-Codestral-7B-v0.1",
27
]
28

29
30
HYBRID_MODELS = [
    "ai21labs/Jamba-tiny-dev",
31
32
33
    # NOTE: ibm-granite/granite-4.0-tiny-preview are skipped currently as
    # it is not yet available in huggingface transformers
    # "ibm-granite/granite-4.0-tiny-preview",
34
35
36
37
38
    # NOTE: Running Plamo2 in transformers implementation requires to install
    # causal-conv1d package, which is not listed as a test dependency as it's
    # not compatible with pip-compile.
    "pfnet/plamo-2-1b",
    "Zyphra/Zamba2-1.2B-instruct",
39
    "hmellor/tiny-random-BambaForCausalLM",
Shinichi Hemmi's avatar
Shinichi Hemmi committed
40
]
41

Chen Zhang's avatar
Chen Zhang committed
42
43
44
45
V1_SUPPORTED_MODELS = [
    "mistralai/Mamba-Codestral-7B-v0.1",
]

46
47
# Avoid OOM
MAX_NUM_SEQS = 4
Mor Zusman's avatar
Mor Zusman committed
48
49


50
51
52
@pytest.mark.parametrize("model", SSM_MODELS + HYBRID_MODELS)
@pytest.mark.parametrize("max_tokens", [64])
@pytest.mark.parametrize("num_logprobs", [5])
Mor Zusman's avatar
Mor Zusman committed
53
54
55
56
def test_models(
    hf_runner,
    vllm_runner,
    example_prompts,
Chen Zhang's avatar
Chen Zhang committed
57
    monkeypatch,
Mor Zusman's avatar
Mor Zusman committed
58
59
    model: str,
    max_tokens: int,
60
    num_logprobs: int,
Mor Zusman's avatar
Mor Zusman committed
61
) -> None:
62
    with hf_runner(model) as hf_model:
Chen Zhang's avatar
Chen Zhang committed
63
64
65
66
67
        if model != "mistralai/Mamba-Codestral-7B-v0.1":
            hf_outputs = hf_model.generate_greedy_logprobs_limit(
                example_prompts, max_tokens, num_logprobs)
        else:
            hf_outputs = None
Yu Chin Fabian Lim's avatar
Yu Chin Fabian Lim committed
68

69
    with vllm_runner(model, max_num_seqs=MAX_NUM_SEQS) as vllm_model:
Chen Zhang's avatar
Chen Zhang committed
70
        vllm_v0_outputs = vllm_model.generate_greedy_logprobs(
71
            example_prompts, max_tokens, num_logprobs)
72

Chen Zhang's avatar
Chen Zhang committed
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
    if model in V1_SUPPORTED_MODELS:
        with monkeypatch.context() as m:
            m.setenv("VLLM_USE_V1", "1")
            with vllm_runner(model,
                             max_num_seqs=MAX_NUM_SEQS,
                             enforce_eager=True,
                             enable_prefix_caching=False) as vllm_model:
                vllm_v1_outputs = vllm_model.generate_greedy_logprobs(
                    example_prompts, max_tokens, num_logprobs)
    else:
        vllm_v1_outputs = None

    if hf_outputs is not None:
        check_logprobs_close(
            outputs_0_lst=hf_outputs,
            outputs_1_lst=vllm_v0_outputs,
            name_0="hf",
            name_1="vllm-v0",
        )

    if model in V1_SUPPORTED_MODELS:
        ref_outputs = hf_outputs if hf_outputs is not None else vllm_v0_outputs
        check_logprobs_close(
            outputs_0_lst=ref_outputs,
            outputs_1_lst=vllm_v1_outputs,
            name_0="hf" if hf_outputs is not None else "vllm-v0",
            name_1="vllm-v1",
        )
Mor Zusman's avatar
Mor Zusman committed
101
102


103
104
105
@pytest.mark.parametrize("model", SSM_MODELS + HYBRID_MODELS)
@pytest.mark.parametrize("max_tokens", [64])
@pytest.mark.parametrize("num_logprobs", [5])
106
107
108
109
110
def test_batching(
    vllm_runner,
    example_prompts,
    model: str,
    max_tokens: int,
111
    num_logprobs: int,
112
113
) -> None:
    for_loop_outputs = []
114
    with vllm_runner(model, max_num_seqs=MAX_NUM_SEQS) as vllm_model:
115
        for prompt in example_prompts:
116
117
118
119
            single_output, = vllm_model.generate_greedy_logprobs([prompt],
                                                                 max_tokens,
                                                                 num_logprobs)
            for_loop_outputs.append(single_output)
120

121
122
        batched_outputs = vllm_model.generate_greedy_logprobs(
            example_prompts, max_tokens, num_logprobs)
123

124
    check_logprobs_close(
125
126
127
128
129
130
131
        outputs_0_lst=for_loop_outputs,
        outputs_1_lst=batched_outputs,
        name_0="for_loop_vllm",
        name_1="batched_vllm",
    )


132
133
134
135
136
137
138
139
140
141
142
143
144
145
@pytest.mark.parametrize("model", [SSM_MODELS[0], HYBRID_MODELS[0]])
@pytest.mark.parametrize("max_tokens", [32])
@pytest.mark.parametrize("num_logprobs", [5])
@pytest.mark.parametrize("chunked_prefill_token_size", [1, 4, 16])
def test_chunked_prefill(
    vllm_runner,
    example_prompts,
    model: str,
    max_tokens: int,
    num_logprobs: int,
    chunked_prefill_token_size: int,
) -> None:
    max_num_seqs = chunked_prefill_token_size
    max_num_batched_tokens = chunked_prefill_token_size
146
147
148

    with vllm_runner(model,
                     enable_chunked_prefill=True,
149
150
151
152
                     max_num_batched_tokens=max_num_batched_tokens,
                     max_num_seqs=max_num_seqs) as vllm_model:
        chunked = vllm_model.generate_greedy_logprobs(example_prompts,
                                                      max_tokens, num_logprobs)
153

154
155
156
157
158
159
160
    with vllm_runner(model,
                     enable_chunked_prefill=False,
                     max_num_seqs=max_num_seqs) as vllm_model:
        non_chunked = vllm_model.generate_greedy_logprobs(
            example_prompts, max_tokens, num_logprobs)

    check_logprobs_close(
161
162
163
164
165
166
167
        outputs_0_lst=chunked,
        outputs_1_lst=non_chunked,
        name_0="chunked",
        name_1="non_chunked",
    )


168
169
170
@pytest.mark.parametrize("model", [SSM_MODELS[0], HYBRID_MODELS[0]])
@pytest.mark.parametrize("max_tokens", [10])
def test_chunked_prefill_with_parallel_sampling(
171
172
173
174
175
    vllm_runner,
    example_prompts,
    model: str,
    max_tokens: int,
) -> None:
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
    """
    Tests chunked prefill in conjunction with n > 1. 
    
    In this case, prefill is populated with decoding tokens and
    we test that it doesn't fail.

    This test might fail if cache is not allocated correctly for n > 1
    decoding steps inside a chunked prefill forward pass
    (where we have both prefill and decode together)
    """
    sampling_params = SamplingParams(n=3,
                                     temperature=1,
                                     seed=0,
                                     max_tokens=max_tokens)
    with vllm_runner(
            model,
            enable_chunked_prefill=True,
            # forces prefill chunks with decoding
            max_num_batched_tokens=MAX_NUM_SEQS * 3,
            max_num_seqs=MAX_NUM_SEQS,
    ) as vllm_model:
        vllm_model.generate(example_prompts, sampling_params)
198
199


200
@pytest.mark.parametrize("model", [SSM_MODELS[0], HYBRID_MODELS[0]])
201
202
203
204
205
206
207
@pytest.mark.parametrize("max_tokens", [20])
def test_mamba_cache_cg_padding(
    vllm_runner,
    example_prompts,
    model: str,
    max_tokens: int,
) -> None:
208
209
210
211
212
    """
    This test is for verifying that mamba cache is padded to CG captured
    batch size. If it's not, a torch RuntimeError will be raised because
    tensor dimensions aren't compatible.
    """
Shinichi Hemmi's avatar
Shinichi Hemmi committed
213
214
    vllm_config = EngineArgs(model=model,
                             trust_remote_code=True).create_engine_config()
215
    while len(example_prompts) == vllm_config.pad_for_cudagraph(
216
            len(example_prompts)):
217
218
219
        example_prompts.append(example_prompts[0])

    try:
220
        with vllm_runner(model) as vllm_model:
221
222
223
224
225
226
227
228
            vllm_model.generate_greedy(example_prompts, max_tokens)
    except RuntimeError:
        pytest.fail(
            "Couldn't run batch size which is not equal to a Cuda Graph "
            "captured batch size. "
            "Could be related to mamba cache not padded correctly")


229
@pytest.mark.parametrize("model", [SSM_MODELS[0], HYBRID_MODELS[0]])
230
231
232
233
234
235
236
@pytest.mark.parametrize("max_tokens", [20])
def test_models_preemption_recompute(
    vllm_runner,
    example_prompts,
    model: str,
    max_tokens: int,
) -> None:
237
238
239
240
241
242
    """
    Tests that outputs are identical with and w/o preemptions (recompute).
    """
    with vllm_runner(model, max_num_seqs=MAX_NUM_SEQS) as vllm_model:
        scheduler = vllm_model.model.llm_engine.scheduler[0]
        scheduler.ENABLE_ARTIFICIAL_PREEMPT = True
243
244
245
        preempt_vllm_outputs = vllm_model.generate_greedy(
            example_prompts, max_tokens)

246
        scheduler.ENABLE_ARTIFICIAL_PREEMPT = False
247
248
249
250
251
252
253
254
255
256
        vllm_outputs = vllm_model.generate_greedy(example_prompts, max_tokens)

    check_outputs_equal(
        outputs_0_lst=preempt_vllm_outputs,
        outputs_1_lst=vllm_outputs,
        name_0="vllm_preepmtions",
        name_1="vllm",
    )


257
@pytest.mark.parametrize("model", [SSM_MODELS[0], HYBRID_MODELS[0]])
258
259
260
def test_fail_upon_inc_requests_and_finished_requests_lt_available_blocks(
    vllm_runner,
    example_prompts,
261
    model: str,
262
) -> None:
263
264
265
266
267
268
269
270
271
    """
    This test is for verifying that the hybrid inner state management doesn't
    collapse in case where the number of incoming requests and
    finished_requests_ids is larger than the maximum mamba block capacity.

    This could generally happen due to the fact that hybrid does support
    statelessness mechanism where it can cleanup new incoming requests in
    a single step.
    """
272
    try:
273
        with vllm_runner(model, max_num_seqs=MAX_NUM_SEQS) as vllm_model:
274
275
            vllm_model.generate_greedy([example_prompts[0]] * 100, 10)
    except ValueError:
Yu Chin Fabian Lim's avatar
Yu Chin Fabian Lim committed
276
        pytest.fail("Hybrid inner state wasn't cleaned up properly between"
277
278
279
                    "steps finished requests registered unnecessarily ")


280
@pytest.mark.parametrize("model", [SSM_MODELS[0], HYBRID_MODELS[0]])
Mor Zusman's avatar
Mor Zusman committed
281
282
283
def test_state_cleanup(
    vllm_runner,
    example_prompts,
284
    model: str,
Mor Zusman's avatar
Mor Zusman committed
285
) -> None:
286
287
288
289
290
291
    """ 
    This test is for verifying that the Hybrid state is cleaned up between
    steps.
    
    If its not cleaned, an error would be expected.
    """
Mor Zusman's avatar
Mor Zusman committed
292
    try:
293
        with vllm_runner(model, max_num_seqs=MAX_NUM_SEQS) as vllm_model:
Mor Zusman's avatar
Mor Zusman committed
294
295
296
            for _ in range(10):
                vllm_model.generate_greedy([example_prompts[0]] * 100, 1)
    except ValueError:
Yu Chin Fabian Lim's avatar
Yu Chin Fabian Lim committed
297
        pytest.fail("Hybrid inner state wasn't cleaned up between states, "
Mor Zusman's avatar
Mor Zusman committed
298
299
300
                    "could be related to finished_requests_ids")


301
302
303
@pytest.mark.parametrize("model", [SSM_MODELS[0], HYBRID_MODELS[0]])
@pytest.mark.parametrize("max_tokens", [64])
def test_multistep_correctness(
304
305
    vllm_runner,
    example_prompts,
306
307
    model: str,
    max_tokens: int,
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
) -> None:
    with vllm_runner(model, num_scheduler_steps=8,
                     max_num_seqs=2) as vllm_model:
        vllm_outputs_multistep = vllm_model.generate_greedy(
            example_prompts, max_tokens)

    with vllm_runner(model, num_scheduler_steps=1,
                     max_num_seqs=2) as vllm_model:
        vllm_outputs_single_step = vllm_model.generate_greedy(
            example_prompts, max_tokens)

    check_outputs_equal(
        outputs_0_lst=vllm_outputs_multistep,
        outputs_1_lst=vllm_outputs_single_step,
        name_0="vllm_outputs_multistep",
        name_1="vllm_outputs_single_step",
    )


327
@multi_gpu_test(num_gpus=2)
328
@pytest.mark.parametrize("model", [SSM_MODELS[0], HYBRID_MODELS[0]])
329
@pytest.mark.parametrize("max_tokens", [64])
330
331
@pytest.mark.parametrize("num_logprobs", [5])
def test_distributed_correctness(
332
333
334
335
    vllm_runner,
    example_prompts,
    model: str,
    max_tokens: int,
336
    num_logprobs: int,
337
) -> None:
338
    with vllm_runner(model, tensor_parallel_size=1,
339
                     max_num_seqs=2) as vllm_model:
340
341
        vllm_outputs_tp_1 = vllm_model.generate_greedy_logprobs(
            example_prompts, max_tokens, num_logprobs)
342

343
    with vllm_runner(model, tensor_parallel_size=2,
344
                     max_num_seqs=2) as vllm_model:
345
346
        vllm_outputs_tp_2 = vllm_model.generate_greedy_logprobs(
            example_prompts, max_tokens, num_logprobs)
347

348
    check_logprobs_close(
349
350
351
352
353
        outputs_0_lst=vllm_outputs_tp_1,
        outputs_1_lst=vllm_outputs_tp_2,
        name_0="vllm_tp_1",
        name_1="vllm_tp_2",
    )