test_hybrid.py 17.2 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
import pytest
5
import os
Mor Zusman's avatar
Mor Zusman committed
6

7
from tests.models.registry import HF_EXAMPLE_MODELS
8
from tests.utils import multi_gpu_test
9
from vllm.engine.arg_utils import EngineArgs
10
from vllm.sampling_params import SamplingParams
11

12
from ....utils import models_path_prefix
13
14
from ...utils import check_logprobs_close, check_outputs_equal

15
16
17
# Mark all tests as hybrid
pytestmark = pytest.mark.hybrid_model

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 = [
23
24
25
    os.path.join(models_path_prefix, "state-spaces/mamba-130m-hf"),
    os.path.join(models_path_prefix, "tiiuae/falcon-mamba-tiny-dev"),
    os.path.join(models_path_prefix, "yujiepan/mamba2-codestral-v0.1-tiny-random"),
26
]
27

28
HYBRID_MODELS = [
29
    os.path.join(models_path_prefix, "ai21labs/Jamba-tiny-dev"),
30
31
    # skipping until vLLM implementation issues are resolved
    # "pfnet/plamo-2-1b",
32
33
34
35
    os.path.join(models_path_prefix, "Zyphra/Zamba2-1.2B-instruct"),
    os.path.join(models_path_prefix, "hmellor/tiny-random-BambaForCausalLM"),
    os.path.join(models_path_prefix, "ibm-granite/granite-4.0-tiny-preview"),
    os.path.join(models_path_prefix, "tiiuae/Falcon-H1-0.5B-Base"),
36
37
38
39
40
41
42
]

HF_UNSUPPORTED_MODELS = [
    # The HF transformers implementation of
    # Mamba2 is buggy for Codestral as it doesn't handle n_groups, so the test
    # doesn't compare vLLM output with HF output.
    # See https://github.com/huggingface/transformers/pull/35943
43
    os.path.join(models_path_prefix, "yujiepan/mamba2-codestral-v0.1-tiny-random"),
44
45
    # transformers 4.55 is still producing garbage for this model
    # TODO(tdoublep): follow-up on transformers side
46
    os.path.join(models_path_prefix, "ibm-granite/granite-4.0-tiny-preview")
Shinichi Hemmi's avatar
Shinichi Hemmi committed
47
]
48

Chen Zhang's avatar
Chen Zhang committed
49
V1_SUPPORTED_MODELS = [
50
51
52
53
54
55
56
    os.path.join(models_path_prefix, "state-spaces/mamba-130m-hf"),
    os.path.join(models_path_prefix, "ai21labs/Jamba-tiny-dev"),
    os.path.join(models_path_prefix, "yujiepan/mamba2-codestral-v0.1-tiny-random"),
    os.path.join(models_path_prefix, "Zyphra/Zamba2-1.2B-instruct"),
    os.path.join(models_path_prefix, "hmellor/tiny-random-BambaForCausalLM"),
    os.path.join(models_path_prefix, "ibm-granite/granite-4.0-tiny-preview"),
    os.path.join(models_path_prefix, "tiiuae/Falcon-H1-0.5B-Base"),
Shinichi Hemmi's avatar
Shinichi Hemmi committed
57
]
58

59

60
61
# Avoid OOM
MAX_NUM_SEQS = 4
Mor Zusman's avatar
Mor Zusman committed
62

63
64
65
# Once we add support for FCG in Mamba1, this list will be removed and tests
# all test cases will use enforce_eager=False
ENFORCE_EAGER_MODELS_V1 = [
66
67
    os.path.join(models_path_prefix, "state-spaces/mamba-130m-hf"),
    os.path.join(models_path_prefix, "ai21labs/Jamba-tiny-dev"),
68
69
]

Mor Zusman's avatar
Mor Zusman committed
70

71
72
73
@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
74
75
76
77
def test_models(
    hf_runner,
    vllm_runner,
    example_prompts,
Chen Zhang's avatar
Chen Zhang committed
78
    monkeypatch,
Mor Zusman's avatar
Mor Zusman committed
79
80
    model: str,
    max_tokens: int,
81
    num_logprobs: int,
Mor Zusman's avatar
Mor Zusman committed
82
) -> None:
83
84
85
86

    try:
        model_info = HF_EXAMPLE_MODELS.find_hf_info(model)
        model_info.check_available_online(on_fail="skip")
87
88
        hf_version_check = model_info.check_transformers_version(
            on_fail="return")
89
    except ValueError:
90
91
92
93
        hf_version_check = None

    if hf_version_check is not None:
        print(f"Skipping transformers comparison because: {hf_version_check}")
94

95
    with hf_runner(model) as hf_model:
96
        if model not in HF_UNSUPPORTED_MODELS and hf_version_check is None:
Chen Zhang's avatar
Chen Zhang committed
97
98
99
100
            hf_outputs = hf_model.generate_greedy_logprobs_limit(
                example_prompts, max_tokens, num_logprobs)
        else:
            hf_outputs = None
Mor Zusman's avatar
Mor Zusman committed
101

102
    with vllm_runner(model, max_num_seqs=MAX_NUM_SEQS) as vllm_model:
Chen Zhang's avatar
Chen Zhang committed
103
        vllm_v0_outputs = vllm_model.generate_greedy_logprobs(
104
            example_prompts, max_tokens, num_logprobs)
105

Chen Zhang's avatar
Chen Zhang committed
106
    if model in V1_SUPPORTED_MODELS:
107
        enforce_eager = False
Chen Zhang's avatar
Chen Zhang committed
108
109
        with monkeypatch.context() as m:
            m.setenv("VLLM_USE_V1", "1")
110
111
112
            if model in HYBRID_MODELS:
                # required due to reorder_batch behaviour
                m.setenv("VLLM_ATTENTION_BACKEND", "FLASHINFER")
113
114
115
116

            if model in ENFORCE_EAGER_MODELS_V1:
                enforce_eager = True

Chen Zhang's avatar
Chen Zhang committed
117
118
            with vllm_runner(model,
                             max_num_seqs=MAX_NUM_SEQS,
119
                             enforce_eager=enforce_eager,
120
                             enable_prefix_caching=False) as vllm_model:
Chen Zhang's avatar
Chen Zhang committed
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
                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
142
143


144
145
146
@pytest.mark.parametrize("model", SSM_MODELS + HYBRID_MODELS)
@pytest.mark.parametrize("max_tokens", [64])
@pytest.mark.parametrize("num_logprobs", [5])
147
148
149
150
151
def test_batching(
    vllm_runner,
    example_prompts,
    model: str,
    max_tokens: int,
152
    num_logprobs: int,
153
) -> None:
154
155
156
157
158
159
160
161

    try:
        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")
    except ValueError:
        pass

162
    for_loop_outputs = []
163
    with vllm_runner(model, max_num_seqs=MAX_NUM_SEQS) as vllm_model:
164
        for prompt in example_prompts:
165
166
167
168
            single_output, = vllm_model.generate_greedy_logprobs([prompt],
                                                                 max_tokens,
                                                                 num_logprobs)
            for_loop_outputs.append(single_output)
169

170
171
        batched_outputs = vllm_model.generate_greedy_logprobs(
            example_prompts, max_tokens, num_logprobs)
172

173
    check_logprobs_close(
174
175
176
177
178
179
180
        outputs_0_lst=for_loop_outputs,
        outputs_1_lst=batched_outputs,
        name_0="for_loop_vllm",
        name_1="batched_vllm",
    )


181
182
183
184
185
186
187
188
189
190
191
192
193
194
@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
195
196
197

    with vllm_runner(model,
                     enable_chunked_prefill=True,
198
199
200
201
                     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)
202

203
204
205
206
207
208
209
    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(
210
211
212
213
214
215
216
        outputs_0_lst=chunked,
        outputs_1_lst=non_chunked,
        name_0="chunked",
        name_1="non_chunked",
    )


217
218
219
@pytest.mark.parametrize("model", [SSM_MODELS[0], HYBRID_MODELS[0]])
@pytest.mark.parametrize("max_tokens", [10])
def test_chunked_prefill_with_parallel_sampling(
220
221
222
223
224
    vllm_runner,
    example_prompts,
    model: str,
    max_tokens: int,
) -> None:
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
    """
    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)
247
248


249
@pytest.mark.parametrize("model", [SSM_MODELS[0], HYBRID_MODELS[0]])
250
251
252
253
254
255
256
@pytest.mark.parametrize("max_tokens", [20])
def test_mamba_cache_cg_padding(
    vllm_runner,
    example_prompts,
    model: str,
    max_tokens: int,
) -> None:
257
258
259
260
261
    """
    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
262
263
    vllm_config = EngineArgs(model=model,
                             trust_remote_code=True).create_engine_config()
264
    while len(example_prompts) == vllm_config.pad_for_cudagraph(
265
            len(example_prompts)):
266
267
268
        example_prompts.append(example_prompts[0])

    try:
269
        with vllm_runner(model) as vllm_model:
270
271
272
273
274
275
276
277
            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")


278
@pytest.mark.parametrize("model", [SSM_MODELS[0], HYBRID_MODELS[0]])
279
280
281
282
283
284
285
@pytest.mark.parametrize("max_tokens", [20])
def test_models_preemption_recompute(
    vllm_runner,
    example_prompts,
    model: str,
    max_tokens: int,
) -> None:
286
287
288
289
    """
    Tests that outputs are identical with and w/o preemptions (recompute).
    """
    with vllm_runner(model, max_num_seqs=MAX_NUM_SEQS) as vllm_model:
290
        scheduler = vllm_model.llm.llm_engine.scheduler[0]
291
        scheduler.ENABLE_ARTIFICIAL_PREEMPT = True
292
293
294
        preempt_vllm_outputs = vllm_model.generate_greedy(
            example_prompts, max_tokens)

295
        scheduler.ENABLE_ARTIFICIAL_PREEMPT = False
296
297
298
299
300
301
302
303
304
305
        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",
    )


306
@pytest.mark.parametrize("model", [SSM_MODELS[0], HYBRID_MODELS[0]])
307
308
309
def test_fail_upon_inc_requests_and_finished_requests_lt_available_blocks(
    vllm_runner,
    example_prompts,
310
    model: str,
311
) -> None:
312
313
314
315
316
317
318
319
320
    """
    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.
    """
321
    try:
322
        with vllm_runner(model, max_num_seqs=MAX_NUM_SEQS) as vllm_model:
323
324
            vllm_model.generate_greedy([example_prompts[0]] * 100, 10)
    except ValueError:
Yu Chin Fabian Lim's avatar
Yu Chin Fabian Lim committed
325
        pytest.fail("Hybrid inner state wasn't cleaned up properly between"
326
327
328
                    "steps finished requests registered unnecessarily ")


329
@pytest.mark.parametrize("model", [SSM_MODELS[0], HYBRID_MODELS[0]])
Mor Zusman's avatar
Mor Zusman committed
330
331
332
def test_state_cleanup(
    vllm_runner,
    example_prompts,
333
    model: str,
Mor Zusman's avatar
Mor Zusman committed
334
) -> None:
335
336
337
338
339
340
    """ 
    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
341
    try:
342
        with vllm_runner(model, max_num_seqs=MAX_NUM_SEQS) as vllm_model:
Mor Zusman's avatar
Mor Zusman committed
343
344
345
            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
346
        pytest.fail("Hybrid inner state wasn't cleaned up between states, "
Mor Zusman's avatar
Mor Zusman committed
347
348
349
                    "could be related to finished_requests_ids")


350
@multi_gpu_test(num_gpus=2)
351
352
@pytest.mark.parametrize("model", [SSM_MODELS[0], HYBRID_MODELS[0]])
@pytest.mark.parametrize("max_tokens", [64])
353
354
@pytest.mark.parametrize("num_logprobs", [5])
def test_distributed_correctness(
Mor Zusman's avatar
Mor Zusman committed
355
    vllm_runner,
356
    example_prompts,
357
358
    model: str,
    max_tokens: int,
359
    num_logprobs: int,
Mor Zusman's avatar
Mor Zusman committed
360
) -> None:
361
    with vllm_runner(model, tensor_parallel_size=1,
362
                     max_num_seqs=2) as vllm_model:
363
364
        vllm_outputs_tp_1 = vllm_model.generate_greedy_logprobs(
            example_prompts, max_tokens, num_logprobs)
365

366
    with vllm_runner(model, tensor_parallel_size=2,
367
                     max_num_seqs=2) as vllm_model:
368
369
        vllm_outputs_tp_2 = vllm_model.generate_greedy_logprobs(
            example_prompts, max_tokens, num_logprobs)
370

371
    check_logprobs_close(
372
373
374
375
        outputs_0_lst=vllm_outputs_tp_1,
        outputs_1_lst=vllm_outputs_tp_2,
        name_0="vllm_tp_1",
        name_1="vllm_tp_2",
376
377
378
    )


379
@pytest.mark.parametrize("model", ["Zyphra/Zamba2-1.2B-instruct"])
380
@pytest.mark.parametrize("max_tokens", [64])
381
@pytest.mark.parametrize("num_logprobs", [5])
382
383
def test_full_cuda_graph(
    hf_runner,
384
385
    vllm_runner,
    example_prompts,
386
    monkeypatch,
387
388
    model: str,
    max_tokens: int,
389
    num_logprobs: int,
390
) -> None:
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407

    try:
        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")
    except ValueError:
        pass

    with hf_runner(model) as hf_model:
        if model not in HF_UNSUPPORTED_MODELS:
            hf_outputs = hf_model.generate_greedy_logprobs_limit(
                example_prompts, max_tokens, num_logprobs)
        else:
            hf_outputs = None

    with vllm_runner(model, max_num_seqs=MAX_NUM_SEQS) as vllm_model:
        vllm_v0_outputs = vllm_model.generate_greedy_logprobs(
408
            example_prompts, max_tokens, num_logprobs)
409

410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
    with monkeypatch.context() as m:
        m.setenv("VLLM_USE_V1", "1")
        if model in HYBRID_MODELS:
            # required due to reorder_batch behaviour
            m.setenv("VLLM_ATTENTION_BACKEND", "FLASHINFER")
        with vllm_runner(model,
                         max_num_seqs=MAX_NUM_SEQS,
                         compilation_config={'full_cuda_graph': True},
                         enable_prefix_caching=False) as vllm_model:
            vllm_v1_outputs = vllm_model.generate_greedy_logprobs(
                example_prompts, max_tokens, num_logprobs)

    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",
        )

    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",
    )
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469


@pytest.mark.parametrize("model", ["Zyphra/Zamba2-1.2B-instruct"])
@pytest.mark.parametrize("max_tokens", [64])
@pytest.mark.parametrize("num_logprobs", [5])
def test_fp32_state(
    hf_runner,
    vllm_runner,
    example_prompts,
    monkeypatch,
    model: str,
    max_tokens: int,
    num_logprobs: int,
) -> None:

    try:
        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")
    except ValueError:
        pass

    with hf_runner(model) as hf_model:
        if model not in HF_UNSUPPORTED_MODELS:
            hf_outputs = hf_model.generate_greedy_logprobs_limit(
                example_prompts, max_tokens, num_logprobs)
        else:
            hf_outputs = None

    with vllm_runner(model,
                     max_num_seqs=MAX_NUM_SEQS,
                     mamba_ssm_cache_dtype="float32") as vllm_model:
        vllm_v0_outputs = vllm_model.generate_greedy_logprobs(
470
            example_prompts, max_tokens, num_logprobs)
471

472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
    with monkeypatch.context() as m:
        m.setenv("VLLM_USE_V1", "1")
        if model in HYBRID_MODELS:
            # required due to reorder_batch behaviour
            m.setenv("VLLM_ATTENTION_BACKEND", "FLASHINFER")
        with vllm_runner(model,
                         max_num_seqs=MAX_NUM_SEQS,
                         mamba_ssm_cache_dtype="float32",
                         enable_prefix_caching=False) as vllm_model:
            vllm_v1_outputs = vllm_model.generate_greedy_logprobs(
                example_prompts, max_tokens, num_logprobs)

    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",
        )

    ref_outputs = hf_outputs if hf_outputs is not None else vllm_v0_outputs
493
    check_logprobs_close(
494
495
496
497
        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",
498
    )