test_hybrid.py 12.6 KB
Newer Older
1
2
# SPDX-License-Identifier: Apache-2.0

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

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

9
10
from ...utils import check_outputs_equal

Yu Chin Fabian Lim's avatar
Yu Chin Fabian Lim committed
11
# This test is for the hybrid models
Shinichi Hemmi's avatar
Shinichi Hemmi committed
12
13
14
15
MODELS = [
    "ai21labs/Jamba-tiny-dev", "Zyphra/Zamba2-1.2B-instruct",
    "pfnet/plamo-2-1b"
]
16
17
# Bamba at Fp32 is too big for the CI (L4 GPU).
# MODELS = ["ai21labs/Jamba-tiny-dev", "ibm-ai-platform/Bamba-9B"]
Shinichi Hemmi's avatar
Shinichi Hemmi committed
18
19
20
# 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.
Mor Zusman's avatar
Mor Zusman committed
21
22
23


@pytest.mark.parametrize("model", MODELS)
24
25
@pytest.mark.parametrize("dtype", ["float"])
@pytest.mark.parametrize("max_tokens", [96])
Mor Zusman's avatar
Mor Zusman committed
26
27
28
29
30
31
32
33
def test_models(
    hf_runner,
    vllm_runner,
    example_prompts,
    model: str,
    dtype: str,
    max_tokens: int,
) -> None:
Yu Chin Fabian Lim's avatar
Yu Chin Fabian Lim committed
34
    # numeric error produces different generation
35
    if "Bamba" in model:
Yu Chin Fabian Lim's avatar
Yu Chin Fabian Lim committed
36
37
        example_prompts.pop(3)

Shinichi Hemmi's avatar
Shinichi Hemmi committed
38
    with hf_runner(model, dtype=dtype) as hf_model:
Mor Zusman's avatar
Mor Zusman committed
39
40
41
42
        hf_outputs = hf_model.generate_greedy(example_prompts, max_tokens)

    with vllm_runner(model, dtype=dtype) as vllm_model:
        vllm_outputs = vllm_model.generate_greedy(example_prompts, max_tokens)
43

Mor Zusman's avatar
Mor Zusman committed
44
45
46
47
48
49
50
51
52
    for i in range(len(example_prompts)):
        hf_output_ids, hf_output_str = hf_outputs[i]
        vllm_output_ids, vllm_output_str = vllm_outputs[i]
        assert hf_output_str == vllm_output_str, (
            f"Test{i}:\nHF: {hf_output_str!r}\nvLLM: {vllm_output_str!r}")
        assert hf_output_ids == vllm_output_ids, (
            f"Test{i}:\nHF: {hf_output_ids}\nvLLM: {vllm_output_ids}")


53
@pytest.mark.parametrize("model", MODELS)
54
55
@pytest.mark.parametrize("dtype", ["float"])
@pytest.mark.parametrize("max_tokens", [96])
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
def test_batching(
    vllm_runner,
    example_prompts,
    model: str,
    dtype: str,
    max_tokens: int,
) -> None:
    # To pass the small model tests, we need full precision.
    for_loop_outputs = []
    with vllm_runner(model, dtype=dtype) as vllm_model:
        for prompt in example_prompts:
            for_loop_outputs.append(
                vllm_model.generate_greedy([prompt], max_tokens)[0])

        batched_outputs = vllm_model.generate_greedy(example_prompts,
                                                     max_tokens)

    check_outputs_equal(
        outputs_0_lst=for_loop_outputs,
        outputs_1_lst=batched_outputs,
        name_0="for_loop_vllm",
        name_1="batched_vllm",
    )


81
82
83
84
85
86
87
88
89
90
91
92
@pytest.mark.parametrize("model", MODELS)
@pytest.mark.parametrize("dtype", ["float16"])
@pytest.mark.parametrize("max_tokens", [10])
def test_mamba_prefill_chunking_with_parallel_sampling(
        hf_runner, vllm_runner, example_prompts, model: str, dtype: str,
        max_tokens: int) -> None:
    # Tests prefill chunking 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 prefills
    # and decoding together )
Shinichi Hemmi's avatar
Shinichi Hemmi committed
93
94
95
96

    if 'plamo-2' in model:
        dtype = "float"  # use a different dtype for plamo

97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
    sampling_params = SamplingParams(n=3,
                                     temperature=1,
                                     seed=0,
                                     max_tokens=max_tokens)
    with vllm_runner(
            model,
            dtype=dtype,
            enable_chunked_prefill=True,
            max_num_batched_tokens=30,
            max_num_seqs=10  # forces prefill chunks with decoding
    ) as vllm_model:
        vllm_model.generate(example_prompts, sampling_params)


@pytest.mark.parametrize("model", MODELS)
@pytest.mark.parametrize("dtype", ["bfloat16"])
Yu Chin Fabian Lim's avatar
Yu Chin Fabian Lim committed
113
@pytest.mark.parametrize("max_tokens", [7])
114
115
116
def test_mamba_prefill_chunking(hf_runner, vllm_runner, example_prompts,
                                model: str, dtype: str,
                                max_tokens: int) -> None:
117
    # numeric error during prefill chunking produces different generation
118
    # compared to w/o prefill chunking for those examples, removed them for now
119
    if "Jamba" in model:
Yu Chin Fabian Lim's avatar
Yu Chin Fabian Lim committed
120
121
122
        example_prompts.pop(7)
        example_prompts.pop(2)
        example_prompts.pop(1)
123
    elif "Bamba" in model:
Yu Chin Fabian Lim's avatar
Yu Chin Fabian Lim committed
124
125
126
127
        example_prompts.pop(6)
        example_prompts.pop(3)
        example_prompts.pop(2)
        dtype = "half"  # use a different dtype for Bamba
Shinichi Hemmi's avatar
Shinichi Hemmi committed
128

129
130
131
    elif "Zamba2" in model:
        example_prompts.pop(7)
        dtype = "half"
Shinichi Hemmi's avatar
Shinichi Hemmi committed
132
133
    elif "plamo-2-1b" in model:
        example_prompts.pop(7)
134

Shinichi Hemmi's avatar
Shinichi Hemmi committed
135
    with hf_runner(model, dtype=dtype) as hf_model:
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
        non_chunked = hf_model.generate_greedy(example_prompts, max_tokens)

    with vllm_runner(model,
                     dtype=dtype,
                     enable_chunked_prefill=True,
                     max_num_batched_tokens=5,
                     max_num_seqs=2) as vllm_model:
        chunked = vllm_model.generate_greedy(example_prompts,
                                             max_tokens=max_tokens)

    check_outputs_equal(
        outputs_0_lst=chunked,
        outputs_1_lst=non_chunked,
        name_0="chunked",
        name_1="non_chunked",
    )


@pytest.mark.parametrize("model", MODELS)
Yu Chin Fabian Lim's avatar
Yu Chin Fabian Lim committed
155
@pytest.mark.parametrize("dtype", ["float"])
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
@pytest.mark.parametrize("max_tokens", [15])
def test_parallel_sampling(
    vllm_runner,
    example_prompts,
    model: str,
    dtype: str,
    max_tokens: int,
) -> None:

    with vllm_runner(model, dtype=dtype) as vllm_model:
        for_loop_outputs = []
        for _ in range(10):
            for_loop_outputs.append(
                # using example_prompts index 1 instead of 0 since with 0 the
                # logprobs get really close and the test doesn't pass
                vllm_model.generate_greedy([example_prompts[1]], max_tokens)
                [0])
        sampling_params = SamplingParams(n=10,
                                         temperature=0.001,
                                         seed=0,
                                         max_tokens=max_tokens)
        n_lt_1_outputs = vllm_model.generate([example_prompts[1]],
                                             sampling_params)
    token_ids, texts = n_lt_1_outputs[0]
    n_lt_1_outputs = [(token_id, text)
                      for token_id, text in zip(token_ids, texts)]

    check_outputs_equal(
        outputs_0_lst=n_lt_1_outputs,
        outputs_1_lst=for_loop_outputs,
        name_0="vllm_n_lt_1_outputs",
        name_1="vllm",
    )


191
@pytest.mark.skip(reason="RE-ENABLE: test is currently failing on main.")
192
193
194
195
196
197
198
199
200
201
202
203
204
@pytest.mark.parametrize("model", MODELS)
@pytest.mark.parametrize("dtype", ["bfloat16"])
@pytest.mark.parametrize("max_tokens", [20])
def test_mamba_cache_cg_padding(
    vllm_runner,
    example_prompts,
    model: str,
    dtype: str,
    max_tokens: int,
) -> None:
    # 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
205
206
    vllm_config = EngineArgs(model=model,
                             trust_remote_code=True).create_engine_config()
207
    while len(example_prompts) == vllm_config.pad_for_cudagraph(
208
            len(example_prompts)):
209
210
211
212
213
214
215
216
217
218
219
220
        example_prompts.append(example_prompts[0])

    try:
        with vllm_runner(model, dtype=dtype) as vllm_model:
            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")


221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
@pytest.mark.parametrize("model", MODELS)
@pytest.mark.parametrize("dtype", ["float"])
@pytest.mark.parametrize("max_tokens", [20])
def test_models_preemption_recompute(
    hf_runner,
    vllm_runner,
    example_prompts,
    model: str,
    dtype: str,
    max_tokens: int,
) -> None:
    # Tests that outputs are identical with and w/o preemtions (recompute)
    assert dtype == "float"

    with vllm_runner(model, dtype=dtype) as vllm_model:
        vllm_model.model.llm_engine.scheduler[
            0].ENABLE_ARTIFICIAL_PREEMPT = True
        preempt_vllm_outputs = vllm_model.generate_greedy(
            example_prompts, max_tokens)

        vllm_model.model.llm_engine.scheduler[
            0].ENABLE_ARTIFICIAL_PREEMPT = False
        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",
    )


@pytest.mark.parametrize("model", MODELS)
@pytest.mark.parametrize("dtype", ["float"])
def test_fail_upon_inc_requests_and_finished_requests_lt_available_blocks(
    vllm_runner,
    model: str,
    dtype: str,
    example_prompts,
) -> None:
Yu Chin Fabian Lim's avatar
Yu Chin Fabian Lim committed
261
    # This test is for verifying that the hybrid inner state management doesn't
262
263
    # collapse in case where the number of incoming requests and
    # finished_requests_ids is larger than the maximum mamba block capacity.
Yu Chin Fabian Lim's avatar
Yu Chin Fabian Lim committed
264
    # This could generally happen due to the fact that hybrid does support
265
266
267
268
269
270
    # statelessness mechanism where it can cleanup new incoming requests in
    # a single step.
    try:
        with vllm_runner(model, dtype=dtype, max_num_seqs=10) as vllm_model:
            vllm_model.generate_greedy([example_prompts[0]] * 100, 10)
    except ValueError:
Yu Chin Fabian Lim's avatar
Yu Chin Fabian Lim committed
271
        pytest.fail("Hybrid inner state wasn't cleaned up properly between"
272
273
274
                    "steps finished requests registered unnecessarily ")


Mor Zusman's avatar
Mor Zusman committed
275
276
277
278
279
280
281
282
@pytest.mark.parametrize("model", MODELS)
@pytest.mark.parametrize("dtype", ["float"])
def test_state_cleanup(
    vllm_runner,
    model: str,
    dtype: str,
    example_prompts,
) -> None:
Yu Chin Fabian Lim's avatar
Yu Chin Fabian Lim committed
283
    # This test is for verifying that the Hybrid state is cleaned up between
Mor Zusman's avatar
Mor Zusman committed
284
285
286
287
288
289
    # steps, If its not cleaned, an error would be expected.
    try:
        with vllm_runner(model, dtype=dtype) as vllm_model:
            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
290
        pytest.fail("Hybrid inner state wasn't cleaned up between states, "
Mor Zusman's avatar
Mor Zusman committed
291
292
293
                    "could be related to finished_requests_ids")


294
@pytest.mark.skip(reason="RE-ENABLE: test is currently failing on main.")
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
@pytest.mark.parametrize("model", MODELS)
@pytest.mark.parametrize("dtype", ["float"])
def test_multistep(
    vllm_runner,
    model: str,
    dtype: str,
    example_prompts,
) -> None:
    # This test is verifying that multistep works correctly
    #on mamba-like models
    with vllm_runner(model, num_scheduler_steps=8,
                     max_num_seqs=2) as vllm_model:
        vllm_model.generate_greedy([example_prompts[0]] * 10, 1)


310
@pytest.mark.skip(reason="RE-ENABLE: test is currently failing on main.")
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
@pytest.mark.parametrize("model", MODELS)
@pytest.mark.parametrize("dtype", ["float"])
@pytest.mark.parametrize("max_tokens", [64])
def test_multistep_correctness(vllm_runner, model: str, dtype: str,
                               max_tokens: int, example_prompts) -> 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",
    )


334
335
336
337
@multi_gpu_test(num_gpus=2)
@pytest.mark.parametrize("model", MODELS)
@pytest.mark.parametrize("dtype", ["float"])
@pytest.mark.parametrize("max_tokens", [64])
Yu Chin Fabian Lim's avatar
Yu Chin Fabian Lim committed
338
def test_hybrid_distributed_produces_identical_generation(
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
        vllm_runner, model: str, dtype: str, max_tokens: int,
        example_prompts) -> None:

    with vllm_runner(model, dtype=dtype, tensor_parallel_size=2) as vllm_model:
        vllm_outputs_tp_2 = vllm_model.generate_greedy(example_prompts,
                                                       max_tokens)

    with vllm_runner(model, dtype=dtype, tensor_parallel_size=1) as vllm_model:
        vllm_outputs_tp_1 = vllm_model.generate_greedy(example_prompts,
                                                       max_tokens)

    check_outputs_equal(
        outputs_0_lst=vllm_outputs_tp_1,
        outputs_1_lst=vllm_outputs_tp_2,
        name_0="vllm_tp_1",
        name_1="vllm_tp_2",
    )