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

import pytest
4
import os
5
6
7

from vllm import SamplingParams

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


@pytest.mark.parametrize(
    "common_llm_kwargs",
    [{
15
        "model_name": os.path.join(models_path_prefix, "JackFram/llama-68m"),
16
17
18
19
20
21

        # Skip cuda graph recording for fast test.
        "enforce_eager": True,
    }])
@pytest.mark.parametrize("per_test_common_llm_kwargs", [{}])
@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
22
23
@pytest.mark.parametrize("test_llm_kwargs",
                         [{
24
                             "speculative_model": os.path.join(models_path_prefix, "JackFram/llama-160m"),
25
26
                             "num_speculative_tokens": 3,
                             "disable_logprobs_during_spec_decoding": False,
27
                         }, {
28
                             "speculative_model": os.path.join(models_path_prefix, "JackFram/llama-160m"),
29
30
                             "num_speculative_tokens": 3,
                             "disable_logprobs_during_spec_decoding": True,
31
                         }])
32
33
34
35
36
37
38
39
@pytest.mark.parametrize("batch_size", [8])
@pytest.mark.parametrize(
    "output_len",
    [
        # Use smaller output len for fast test.
        7,
    ])
@pytest.mark.parametrize("seed", [1])
40
41
42
43
44
@pytest.mark.parametrize("logprobs", [1, 6])
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,
                           seed: int, logprobs: int):
45
46
    """Verify output logprobs are equal with and without speculative decoding.
    """
47
48
49
50
51
52
53
54
55
56
57
58
59
    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'])
60
61
62
63
64


@pytest.mark.parametrize(
    "common_llm_kwargs",
    [{
65
        "model_name": os.path.join(models_path_prefix, "JackFram/llama-68m"),
66
67
68
69
70
71

        # Skip cuda graph recording for fast test.
        "enforce_eager": True,
    }])
@pytest.mark.parametrize("per_test_common_llm_kwargs", [{}])
@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
72
73
@pytest.mark.parametrize("test_llm_kwargs",
                         [{
74
                             "speculative_model": os.path.join(models_path_prefix, "JackFram/llama-160m"),
75
76
77
                             "num_speculative_tokens": 3,
                             "disable_logprobs_during_spec_decoding": False,
                         }, {
78
                             "speculative_model": os.path.join(models_path_prefix, "JackFram/llama-160m"),
79
80
81
                             "num_speculative_tokens": 6,
                             "disable_logprobs_during_spec_decoding": False,
                         }])
82
83
84
85
86
87
88
89
@pytest.mark.parametrize("batch_size", [8])
@pytest.mark.parametrize(
    "output_len",
    [
        # Use smaller output len for fast test.
        32,
    ])
@pytest.mark.parametrize("seed", [1])
90
91
92
93
94
@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):
95
96
    """Veriy logprob greedy equality with different speculation lens.
    """
97
98
99
100
101
102
103
104
105
106
107
108
    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'])
109
110
111
112
113


@pytest.mark.parametrize(
    "common_llm_kwargs",
    [{
114
        "model_name": os.path.join(models_path_prefix, "JackFram/llama-68m"),
115
116
117
118
119
120
121
122
123

        # 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",
    [{
124
        "speculative_model": os.path.join(models_path_prefix, "JackFram/llama-160m"),
125
        "num_speculative_tokens": 3,
126
        "disable_logprobs_during_spec_decoding": False,
127
128
129
130
131
132
133
134
135
136
137
138
139

        # 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])
140
141
142
143
144
145
@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):
146
147
    """Verify logprobs greedy equality when some sequences skip speculation.
    """
148
149
150
151
152
153
154
155
156
157
158
159
    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'])
160
161
162
163
164


@pytest.mark.parametrize(
    "common_llm_kwargs",
    [{
165
        "model_name": os.path.join(models_path_prefix, "JackFram/llama-68m"),
166
167
168
169
170
171

        # Skip cuda graph recording for fast test.
        "enforce_eager": True,
    }])
@pytest.mark.parametrize("per_test_common_llm_kwargs", [{}])
@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
172
173
@pytest.mark.parametrize("test_llm_kwargs",
                         [{
174
                             "speculative_model": os.path.join(models_path_prefix, "JackFram/llama-160m"),
175
176
177
                             "num_speculative_tokens": 3,
                             "disable_logprobs_during_spec_decoding": False,
                         }])
178
179
180
181
182
183
184
185
@pytest.mark.parametrize("batch_size", [1])
@pytest.mark.parametrize(
    "output_len",
    [
        # Use smaller output len for fast test.
        32,
    ])
@pytest.mark.parametrize("seed", [1])
186
187
188
189
190
@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):
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
    """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(
213
214
        max_tokens=output_len,
        ignore_eos=True,
215
        temperature=temperature,
216
        logprobs=logprobs,
217
218
    )

219
220
221
222
223
224
225
226
    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)
227
228

    num_returned_logprobs = [
229
        len(seq_logprobs) for seq_logprobs in sd_outputs[-1]
230
231
232
233
    ]

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


@pytest.mark.parametrize(
    "common_llm_kwargs",
    [{
241
        "model_name": os.path.join(models_path_prefix, "JackFram/llama-160m"),
242
243
244
245
246
247
248
        # 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",
                         [{
249
                             "speculative_model": os.path.join(models_path_prefix, "JackFram/llama-68m"),
250
251
252
253
                             "num_speculative_tokens": 3,
                             "disable_logprobs_during_spec_decoding": True,
                         }])
@pytest.mark.parametrize("seed", [1])
254
255
256
257
258
259
260
261
262
263
264
265
@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):
266
267
268
    """Check the behavior when logprobs are disabled.
    Token choices should match with the base model.
    """
269
270
271
272
273
274
275
276
277
278
279
280
    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'])