test_logprobs.py 7.46 KB
Newer Older
1
2
# SPDX-License-Identifier: Apache-2.0

3
4
from typing import List

5
6
import pytest
import torch
7
import os
8
9
10

from vllm import SamplingParams

11
from ..conftest import VllmRunner
12
from ..utils import models_path_prefix
13

14
MODELS = [os.path.join(models_path_prefix, "facebook/opt-125m")]
15
16
17


@pytest.mark.parametrize("model", MODELS)
18
@pytest.mark.parametrize("dtype",
19
                         ["half"])  # needed for comparing logprobs with HF
20
@pytest.mark.parametrize("chunked_prefill_token_size", [1, 4, 16, -1])
21
@pytest.mark.parametrize("num_top_logprobs", [0, 6])  # 32000 == vocab_size
22
@pytest.mark.parametrize("detokenize", [True, False])
23
24
25
26
27
def test_get_prompt_logprobs(
    hf_runner,
    vllm_runner,
    model,
    dtype,
28
29
    chunked_prefill_token_size: int,
    num_top_logprobs: int,
30
    detokenize: bool,
31
32
    example_prompts,
):
33
34
35
36
37
38
39
40
    max_num_seqs = 256
    enable_chunked_prefill = False
    max_num_batched_tokens = None
    if chunked_prefill_token_size != -1:
        enable_chunked_prefill = True
        max_num_seqs = min(chunked_prefill_token_size, max_num_seqs)
        max_num_batched_tokens = chunked_prefill_token_size

41
    max_tokens = 5
42
43
44
45
46
    with hf_runner(model, dtype=dtype) as hf_model:
        hf_logprobs = hf_model.generate_greedy_logprobs(
            example_prompts,
            max_tokens=max_tokens,
        )
47

48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
    with vllm_runner(
            model,
            dtype=dtype,
            max_logprobs=num_top_logprobs,
            enable_chunked_prefill=enable_chunked_prefill,
            max_num_batched_tokens=max_num_batched_tokens,
            max_num_seqs=max_num_seqs,
    ) as vllm_model:
        vllm_sampling_params = SamplingParams(max_tokens=max_tokens,
                                              logprobs=num_top_logprobs,
                                              prompt_logprobs=num_top_logprobs,
                                              temperature=0.0,
                                              detokenize=detokenize)
        vllm_results = vllm_model.model.generate(
            example_prompts, sampling_params=vllm_sampling_params)
63
64
65
66
67

    # Test whether logprobs are included in the results.
    for result in vllm_results:
        assert result.prompt_logprobs is not None
        assert result.outputs[0].logprobs is not None
68
69
        assert len(result.outputs[0].logprobs) == max_tokens
        for logprobs in result.outputs[0].logprobs:
70
71
72
73
            # If the output token is not included in the top X
            # logprob, it can return 1 more data
            assert (len(logprobs) == num_top_logprobs
                    or len(logprobs) == num_top_logprobs + 1)
74
        output_text = result.outputs[0].text
75
        output_string_from_most_likely_tokens_lst: List[str] = []
76
77
        for top_logprobs in result.outputs[0].logprobs:
            top_logprob = next(iter(top_logprobs.values()))
78
            output_string_from_most_likely_tokens_lst.append(
79
                top_logprob.decoded_token)
80
81
82

        if detokenize:
            output_string_from_most_likely_tokens = "".join(
83
                output_string_from_most_likely_tokens_lst)
84
85
86
87
88
            assert output_text == output_string_from_most_likely_tokens, (
                "The output text from the top logprob for each token position "
                "should be the same as the output text in the result.")
        else:
            assert output_text == ''
89
90
            assert output_string_from_most_likely_tokens_lst == ([None] *
                                                                 max_tokens)
91

92
93
94
95
96
97
98
99
        # The first prompt logprob is always None
        assert result.prompt_logprobs[0] is None
        for prompt_logprobs in result.prompt_logprobs[1:]:
            # If the prompt token is not included in the top X
            # logprob, it can return 1 more data
            assert (len(prompt_logprobs) == num_top_logprobs
                    or len(prompt_logprobs) == num_top_logprobs + 1)

100
101
102
    # Test whether prompt logprobs are consistent with HF
    for vllm_result, hf_logprob in zip(vllm_results, hf_logprobs):
        # Check prompt logprobs
103
        # The first prompt logprob is always None, so we compare it from 1:.
104
105
106
        vllm_prompt_logprobs = vllm_result.prompt_logprobs[1:]
        for i, vllm_prompt_logprob_dict in enumerate(vllm_prompt_logprobs):
            for token_id, logprob in vllm_prompt_logprob_dict.items():
107
                torch.testing.assert_close(logprob.logprob,
108
109
110
111
                                           hf_logprob[0][i][token_id].item(),
                                           atol=1e-2,
                                           rtol=1e-2)
        vllm_sample_logprobs = vllm_result.outputs[0].logprobs
112
113
114
        for i, top_logprobs in enumerate(vllm_sample_logprobs):
            for token_id, sample_logprob in top_logprobs.items():
                logprob = sample_logprob.logprob
115
116
                torch.testing.assert_close(logprob,
                                           hf_logprob[i][-1][token_id].item(),
117
118
                                           atol=1e-1,
                                           rtol=1e-1)
119
120
121
122
                if detokenize:
                    assert isinstance(sample_logprob.decoded_token, str), (
                        "The token should be decoded by the time it is returned"
                        " to the user.")
123

124
125
126
127
128
129
130
131
132
133
134
    # Test if prompt logprobs are correctly set.
    for vllm_result in vllm_results:
        token_ids = vllm_result.prompt_token_ids
        prompt_logprobs = vllm_result.prompt_logprobs

        # The first token doesn't have logprob.
        assert prompt_logprobs[0] is None

        for token_id, logprob_dict in zip(token_ids[1:], prompt_logprobs[1:]):
            assert token_id in logprob_dict

135
136

def test_max_logprobs():
137
    runner = VllmRunner(os.path.join(models_path_prefix, "facebook/opt-125m"), max_logprobs=1)
138
139
140
141
142
143
144
    vllm_sampling_params = SamplingParams(logprobs=1)
    # should pass
    runner.generate(["Hello world"], sampling_params=vllm_sampling_params)

    bad_sampling_params = SamplingParams(logprobs=2)
    with pytest.raises(ValueError):
        runner.generate(["Hello world"], sampling_params=bad_sampling_params)
zhuwenwen's avatar
zhuwenwen committed
145

146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176

@pytest.mark.parametrize("model", MODELS)
@pytest.mark.parametrize("chunked_prefill_token_size", [1, 4, 16, -1])
@pytest.mark.parametrize("detokenize", [True, False])
def test_none_logprobs(vllm_runner, model, chunked_prefill_token_size: int,
                       detokenize: bool, example_prompts):
    max_num_seqs = 256
    enable_chunked_prefill = False
    max_num_batched_tokens = None
    if chunked_prefill_token_size != -1:
        enable_chunked_prefill = True
        max_num_seqs = min(chunked_prefill_token_size, max_num_seqs)
        max_num_batched_tokens = chunked_prefill_token_size
    max_tokens = 5

    with vllm_runner(
            model,
            enable_chunked_prefill=enable_chunked_prefill,
            max_num_batched_tokens=max_num_batched_tokens,
            max_num_seqs=max_num_seqs,
    ) as vllm_model:
        sampling_params_logprobs_none = SamplingParams(max_tokens=max_tokens,
                                                       logprobs=None,
                                                       temperature=0.0,
                                                       detokenize=detokenize)
        results_logprobs_none = vllm_model.model.generate(
            example_prompts, sampling_params=sampling_params_logprobs_none)

    for i in range(len(results_logprobs_none)):
        assert results_logprobs_none[i].outputs[0].logprobs is None
        assert results_logprobs_none[i].outputs[0].cumulative_logprob is None