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

4
5
6
7
8
9
from itertools import cycle

import pytest

from vllm import SamplingParams

10
from ..utils import maybe_enable_chunked_prefill
11
from .conftest import run_equality_correctness_test
12
13
14
15
16


@pytest.mark.parametrize(
    "common_llm_kwargs",
    [{
17
        "model_name": "JackFram/llama-160m",
18
19

        # Skip cuda graph recording for fast test.
20
        "enforce_eager": True
21
22
23
    }])
@pytest.mark.parametrize("per_test_common_llm_kwargs", [{}])
@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
24
25
26
27
28
29
30
31
32
33
34
35
36
@pytest.mark.parametrize("test_llm_kwargs", [{
    "speculative_config": {
        "model": "JackFram/llama-68m",
        "num_speculative_tokens": 3,
        "disable_logprobs": False,
    },
}, {
    "speculative_config": {
        "model": "JackFram/llama-68m",
        "num_speculative_tokens": 3,
        "disable_logprobs": True,
    },
}])
37
38
39
40
41
42
43
44
@pytest.mark.parametrize("batch_size", [8])
@pytest.mark.parametrize(
    "output_len",
    [
        # Use smaller output len for fast test.
        7,
    ])
@pytest.mark.parametrize("seed", [1])
45
@pytest.mark.parametrize("logprobs", [1, 6])
46
@pytest.mark.parametrize("prefill_chunk_size", [-1, 4, 12])
47
48
49
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,
50
51
52
                           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.
53
    """
54
    maybe_enable_chunked_prefill(prefill_chunk_size, common_llm_kwargs)
55
56
57
58
59
60
61
62
63
64
65
66
67
68
    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["speculative_config"]
        ["disable_logprobs"])
69
70
71
72
73


@pytest.mark.parametrize(
    "common_llm_kwargs",
    [{
74
        "model_name": "JackFram/llama-68m",
75
76
77
78
79
80

        # Skip cuda graph recording for fast test.
        "enforce_eager": True,
    }])
@pytest.mark.parametrize("per_test_common_llm_kwargs", [{}])
@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
81
82
83
84
85
86
87
88
89
90
91
92
93
@pytest.mark.parametrize("test_llm_kwargs", [{
    "speculative_config": {
        "model": "JackFram/llama-160m",
        "num_speculative_tokens": 3,
        "disable_logprobs": False,
    },
}, {
    "speculative_config": {
        "model": "JackFram/llama-160m",
        "num_speculative_tokens": 6,
        "disable_logprobs": False,
    },
}])
94
95
96
97
98
99
100
101
@pytest.mark.parametrize("batch_size", [8])
@pytest.mark.parametrize(
    "output_len",
    [
        # Use smaller output len for fast test.
        32,
    ])
@pytest.mark.parametrize("seed", [1])
102
103
104
105
106
@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):
107
108
    """Veriy logprob greedy equality with different speculation lens.
    """
109
110
111
112
113
114
115
116
117
118
119
120
121
    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["speculative_config"]
        ["disable_logprobs"])
122
123
124
125
126


@pytest.mark.parametrize(
    "common_llm_kwargs",
    [{
127
        "model_name": "JackFram/llama-68m",
128
129
130
131
132
133
134
135
136

        # 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",
    [{
137
138
139
140
141
142
143
144
145
        "speculative_config": {
            "model": "JackFram/llama-160m",
            "num_speculative_tokens": 3,
            "disable_logprobs": False,
            # Artificially limit the draft model max model len; this forces
            # vLLM to skip speculation once the sequences grow beyond 32-k
            # tokens.
            "max_model_len": 32,
        },
146
147
148
149
150
151
152
153
154
    }])
@pytest.mark.parametrize("batch_size", [8])
@pytest.mark.parametrize(
    "output_len",
    [
        # Use smaller output len for fast test.
        32,
    ])
@pytest.mark.parametrize("seed", [1])
155
156
157
158
159
160
@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):
161
162
    """Verify logprobs greedy equality when some sequences skip speculation.
    """
163
164
165
166
167
168
169
170
171
172
173
174
175
    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["speculative_config"]
        ["disable_logprobs"])
176
177
178
179
180


@pytest.mark.parametrize(
    "common_llm_kwargs",
    [{
181
        "model_name": "JackFram/llama-68m",
182
183
184
185
186
187

        # Skip cuda graph recording for fast test.
        "enforce_eager": True,
    }])
@pytest.mark.parametrize("per_test_common_llm_kwargs", [{}])
@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
188
189
190
191
192
193
194
@pytest.mark.parametrize("test_llm_kwargs", [{
    "speculative_config": {
        "model": "JackFram/llama-160m",
        "num_speculative_tokens": 3,
        "disable_logprobs": False,
    },
}])
195
196
197
198
199
200
201
202
@pytest.mark.parametrize("batch_size", [1])
@pytest.mark.parametrize(
    "output_len",
    [
        # Use smaller output len for fast test.
        32,
    ])
@pytest.mark.parametrize("seed", [1])
203
204
205
206
207
@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):
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
    """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(
230
231
        max_tokens=output_len,
        ignore_eos=True,
232
        temperature=temperature,
233
        logprobs=logprobs,
234
235
    )

236
237
238
239
240
241
242
243
    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)
244
245

    num_returned_logprobs = [
246
        len(seq_logprobs) for seq_logprobs in sd_outputs[-1]
247
248
249
250
    ]

    # Assert one of the returned logprobs has > num_logprobs (indicating the
    # sampled token is not in top-k).
251
252
    assert any(
        [num_returned > logprobs for num_returned in num_returned_logprobs])
253
254
255
256
257


@pytest.mark.parametrize(
    "common_llm_kwargs",
    [{
258
        "model_name": "JackFram/llama-160m",
259
260
261
262
263
        # Skip cuda graph recording for fast test.
        "enforce_eager": True,
    }])
@pytest.mark.parametrize("per_test_common_llm_kwargs", [{}])
@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
264
265
266
267
268
269
270
@pytest.mark.parametrize("test_llm_kwargs", [{
    "speculative_config": {
        "model": "JackFram/llama-68m",
        "num_speculative_tokens": 3,
        "disable_logprobs": True,
    },
}])
271
@pytest.mark.parametrize("seed", [1])
272
273
274
275
276
277
278
279
280
281
282
283
@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):
284
285
286
    """Check the behavior when logprobs are disabled.
    Token choices should match with the base model.
    """
287
288
289
290
291
292
293
294
295
296
297
298
299
    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["speculative_config"]
        ["disable_logprobs"])