test_logprobs.py 11.3 KB
Newer Older
1
2
3
from itertools import cycle

import pytest
4
import os
5
6
7

from vllm import SamplingParams

8
from ..utils import maybe_enable_chunked_prefill
9
from .conftest import run_equality_correctness_test
10
from ...utils import models_path_prefix
11
12
13
14
15


@pytest.mark.parametrize(
    "common_llm_kwargs",
    [{
zhuwenwen's avatar
zhuwenwen committed
16
        "model_name": os.path.join(models_path_prefix, "JackFram/llama-160m"),
17
18

        # Skip cuda graph recording for fast test.
19
        "enforce_eager": True
20
21
22
    }])
@pytest.mark.parametrize("per_test_common_llm_kwargs", [{}])
@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
23
24
@pytest.mark.parametrize("test_llm_kwargs",
                         [{
zhuwenwen's avatar
zhuwenwen committed
25
                             "speculative_model": os.path.join(models_path_prefix, "JackFram/llama-68m"),
26
27
                             "num_speculative_tokens": 3,
                             "disable_logprobs_during_spec_decoding": False,
28
                         }, {
zhuwenwen's avatar
zhuwenwen committed
29
                             "speculative_model": os.path.join(models_path_prefix, "JackFram/llama-68m"),
30
31
                             "num_speculative_tokens": 3,
                             "disable_logprobs_during_spec_decoding": True,
32
                         }])
33
34
35
36
37
38
39
40
@pytest.mark.parametrize("batch_size", [8])
@pytest.mark.parametrize(
    "output_len",
    [
        # Use smaller output len for fast test.
        7,
    ])
@pytest.mark.parametrize("seed", [1])
41
@pytest.mark.parametrize("logprobs", [1, 6])
42
@pytest.mark.parametrize("prefill_chunk_size", [-1, 4, 12])
43
44
45
def test_logprobs_equality(vllm_runner, common_llm_kwargs,
                           per_test_common_llm_kwargs, baseline_llm_kwargs,
                           test_llm_kwargs, batch_size: int, output_len: int,
46
47
48
                           seed: int, logprobs: int, prefill_chunk_size: int):
    """Verify output logprobs are equal with and without speculative decoding,
        as well as with and without chunked prefill.
49
    """
50
    maybe_enable_chunked_prefill(prefill_chunk_size, common_llm_kwargs)
51
52
53
54
55
56
57
58
59
60
61
62
63
    run_equality_correctness_test(vllm_runner,
                                  common_llm_kwargs,
                                  per_test_common_llm_kwargs,
                                  baseline_llm_kwargs,
                                  test_llm_kwargs,
                                  batch_size,
                                  output_len,
                                  seed,
                                  temperature=0.0,
                                  logprobs=logprobs,
                                  prompt_logprobs=logprobs,
                                  disable_logprobs=test_llm_kwargs[
                                      'disable_logprobs_during_spec_decoding'])
64
65
66
67
68


@pytest.mark.parametrize(
    "common_llm_kwargs",
    [{
69
        "model_name": os.path.join(models_path_prefix, "JackFram/llama-68m"),
70
71
72
73
74
75

        # Skip cuda graph recording for fast test.
        "enforce_eager": True,
    }])
@pytest.mark.parametrize("per_test_common_llm_kwargs", [{}])
@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
76
77
@pytest.mark.parametrize("test_llm_kwargs",
                         [{
78
                             "speculative_model": os.path.join(models_path_prefix, "JackFram/llama-160m"),
79
80
81
                             "num_speculative_tokens": 3,
                             "disable_logprobs_during_spec_decoding": False,
                         }, {
82
                             "speculative_model": os.path.join(models_path_prefix, "JackFram/llama-160m"),
83
84
85
                             "num_speculative_tokens": 6,
                             "disable_logprobs_during_spec_decoding": False,
                         }])
86
87
88
89
90
91
92
93
@pytest.mark.parametrize("batch_size", [8])
@pytest.mark.parametrize(
    "output_len",
    [
        # Use smaller output len for fast test.
        32,
    ])
@pytest.mark.parametrize("seed", [1])
94
95
96
97
98
@pytest.mark.parametrize("logprobs", [1, 6])
def test_logprobs_different_k(vllm_runner, common_llm_kwargs,
                              per_test_common_llm_kwargs, baseline_llm_kwargs,
                              test_llm_kwargs, batch_size: int,
                              output_len: int, seed: int, logprobs: int):
99
100
    """Veriy logprob greedy equality with different speculation lens.
    """
101
102
103
104
105
106
107
108
109
110
111
112
    run_equality_correctness_test(vllm_runner,
                                  common_llm_kwargs,
                                  per_test_common_llm_kwargs,
                                  baseline_llm_kwargs,
                                  test_llm_kwargs,
                                  batch_size,
                                  output_len,
                                  seed,
                                  temperature=0.0,
                                  logprobs=logprobs,
                                  disable_logprobs=test_llm_kwargs[
                                      'disable_logprobs_during_spec_decoding'])
113
114
115
116
117


@pytest.mark.parametrize(
    "common_llm_kwargs",
    [{
118
        "model_name": os.path.join(models_path_prefix, "JackFram/llama-68m"),
119
120
121
122
123
124
125
126
127

        # Skip cuda graph recording for fast test.
        "enforce_eager": True,
    }])
@pytest.mark.parametrize("per_test_common_llm_kwargs", [{}])
@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
@pytest.mark.parametrize(
    "test_llm_kwargs",
    [{
128
        "speculative_model": os.path.join(models_path_prefix, "JackFram/llama-160m"),
129
        "num_speculative_tokens": 3,
130
        "disable_logprobs_during_spec_decoding": False,
131
132
133
134
135
136
137
138
139
140
141
142
143

        # Artificially limit the draft model max model len; this forces vLLM
        # to skip speculation once the sequences grow beyond 32-k tokens.
        "speculative_max_model_len": 32,
    }])
@pytest.mark.parametrize("batch_size", [8])
@pytest.mark.parametrize(
    "output_len",
    [
        # Use smaller output len for fast test.
        32,
    ])
@pytest.mark.parametrize("seed", [1])
144
145
146
147
148
149
@pytest.mark.parametrize("logprobs", [1])
def test_logprobs_when_skip_speculation(vllm_runner, common_llm_kwargs,
                                        per_test_common_llm_kwargs,
                                        baseline_llm_kwargs, test_llm_kwargs,
                                        batch_size: int, output_len: int,
                                        seed: int, logprobs: int):
150
151
    """Verify logprobs greedy equality when some sequences skip speculation.
    """
152
153
154
155
156
157
158
159
160
161
162
163
    run_equality_correctness_test(vllm_runner,
                                  common_llm_kwargs,
                                  per_test_common_llm_kwargs,
                                  baseline_llm_kwargs,
                                  test_llm_kwargs,
                                  batch_size,
                                  output_len,
                                  seed,
                                  temperature=0.0,
                                  logprobs=logprobs,
                                  disable_logprobs=test_llm_kwargs[
                                      'disable_logprobs_during_spec_decoding'])
164
165
166
167
168


@pytest.mark.parametrize(
    "common_llm_kwargs",
    [{
169
        "model_name": os.path.join(models_path_prefix, "JackFram/llama-68m"),
170
171
172
173
174
175

        # Skip cuda graph recording for fast test.
        "enforce_eager": True,
    }])
@pytest.mark.parametrize("per_test_common_llm_kwargs", [{}])
@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
176
177
@pytest.mark.parametrize("test_llm_kwargs",
                         [{
178
                             "speculative_model": os.path.join(models_path_prefix, "JackFram/llama-160m"),
179
180
181
                             "num_speculative_tokens": 3,
                             "disable_logprobs_during_spec_decoding": False,
                         }])
182
183
184
185
186
187
188
189
@pytest.mark.parametrize("batch_size", [1])
@pytest.mark.parametrize(
    "output_len",
    [
        # Use smaller output len for fast test.
        32,
    ])
@pytest.mark.parametrize("seed", [1])
190
191
192
193
194
@pytest.mark.parametrize("logprobs", [6])
def test_logprobs_temp_1(vllm_runner, common_llm_kwargs,
                         per_test_common_llm_kwargs, baseline_llm_kwargs,
                         test_llm_kwargs, batch_size: int, output_len: int,
                         seed: int, logprobs: int):
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
    """Verify at least one logprob result has num_logprobs+1, which tests the
    case where the sampled token is not in top-k logprobs.

    Ideally, this test should validate equality with non-spec by getting
    logprobs. This is left as future improvement.
    """
    temperature = 1.0

    prompts = [
        "Hello, my name is",
        "The president of the United States is",
        "The capital of France is",
        "The future of AI is",
        "San Francisco is know for its",
        "Facebook was created in 2004 by",
        "Curious George is a",
        "Python 3.11 brings improvements to its",
    ]

    prompts = [prompt for prompt, _ in zip(cycle(prompts), range(batch_size))]

    sampling_params = SamplingParams(
217
218
        max_tokens=output_len,
        ignore_eos=True,
219
        temperature=temperature,
220
        logprobs=logprobs,
221
222
    )

223
224
225
226
227
228
229
230
    sd_args = {
        **common_llm_kwargs,
        **per_test_common_llm_kwargs,
        **test_llm_kwargs,
    }

    with vllm_runner(**sd_args) as vllm_model:
        sd_outputs = vllm_model.generate_w_logprobs(prompts, sampling_params)
231
232

    num_returned_logprobs = [
233
        len(seq_logprobs) for seq_logprobs in sd_outputs[-1]
234
235
236
237
    ]

    # Assert one of the returned logprobs has > num_logprobs (indicating the
    # sampled token is not in top-k).
238
239
    assert any(
        [num_returned > logprobs for num_returned in num_returned_logprobs])
240
241
242
243
244


@pytest.mark.parametrize(
    "common_llm_kwargs",
    [{
245
        "model_name": os.path.join(models_path_prefix, "JackFram/llama-160m"),
246
247
248
249
250
251
252
        # Skip cuda graph recording for fast test.
        "enforce_eager": True,
    }])
@pytest.mark.parametrize("per_test_common_llm_kwargs", [{}])
@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
@pytest.mark.parametrize("test_llm_kwargs",
                         [{
253
                             "speculative_model": os.path.join(models_path_prefix, "JackFram/llama-68m"),
254
255
256
257
                             "num_speculative_tokens": 3,
                             "disable_logprobs_during_spec_decoding": True,
                         }])
@pytest.mark.parametrize("seed", [1])
258
259
260
261
262
263
264
265
266
267
268
269
@pytest.mark.parametrize("batch_size", [4])
@pytest.mark.parametrize(
    "output_len",
    [
        # Use smaller output len for fast test.
        32,
    ])
@pytest.mark.parametrize("logprobs", [0])
def test_logprobs_disabled(vllm_runner, common_llm_kwargs,
                           per_test_common_llm_kwargs, baseline_llm_kwargs,
                           test_llm_kwargs, batch_size: int, output_len: int,
                           seed: int, logprobs: int):
270
271
272
    """Check the behavior when logprobs are disabled.
    Token choices should match with the base model.
    """
273
274
275
276
277
278
279
280
281
282
283
284
    run_equality_correctness_test(vllm_runner,
                                  common_llm_kwargs,
                                  per_test_common_llm_kwargs,
                                  baseline_llm_kwargs,
                                  test_llm_kwargs,
                                  batch_size,
                                  output_len,
                                  seed,
                                  temperature=0.0,
                                  logprobs=logprobs,
                                  disable_logprobs=test_llm_kwargs[
                                      'disable_logprobs_during_spec_decoding'])