test_logprobs.py 5.7 KB
Newer Older
1
2
from typing import List

3
4
5
6
7
import pytest
import torch

from vllm import SamplingParams

8
9
from ..conftest import VllmRunner

10
11
12
13
MODELS = ["facebook/opt-125m"]


@pytest.mark.parametrize("model", MODELS)
14
15
@pytest.mark.parametrize("dtype",
                         ["float"])  # needed for comparing logprobs with HF
16
17
@pytest.mark.parametrize("chunked_prefill_token_size", [1, 4, 16, -1])
@pytest.mark.parametrize("num_top_logprobs", [6])  # 32000 == vocab_size
18
@pytest.mark.parametrize("detokenize", [True, False])
19
20
21
22
23
def test_get_prompt_logprobs(
    hf_runner,
    vllm_runner,
    model,
    dtype,
24
25
    chunked_prefill_token_size: int,
    num_top_logprobs: int,
26
    detokenize: bool,
27
28
    example_prompts,
):
29
30
31
32
33
34
35
36
    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

37
    max_tokens = 5
38
39
40
41
42
    with hf_runner(model, dtype=dtype) as hf_model:
        hf_logprobs = hf_model.generate_greedy_logprobs(
            example_prompts,
            max_tokens=max_tokens,
        )
43

44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
    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)
59
60
61
62
63

    # 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
64
65
66
67
        assert len(result.outputs[0].logprobs) == max_tokens
        for logprobs in result.outputs[0].logprobs:
            assert len(logprobs) == num_top_logprobs
        output_text = result.outputs[0].text
68
        output_string_from_most_likely_tokens_lst: List[str] = []
69
70
        for top_logprobs in result.outputs[0].logprobs:
            top_logprob = next(iter(top_logprobs.values()))
71
            output_string_from_most_likely_tokens_lst.append(
72
                top_logprob.decoded_token)
73
74
75

        if detokenize:
            output_string_from_most_likely_tokens = "".join(
76
                output_string_from_most_likely_tokens_lst)
77
78
79
80
81
            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 == ''
82
83
            assert output_string_from_most_likely_tokens_lst == ([None] *
                                                                 max_tokens)
84

85
86
87
88
89
90
91
92
        # 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)

93
94
95
    # Test whether prompt logprobs are consistent with HF
    for vllm_result, hf_logprob in zip(vllm_results, hf_logprobs):
        # Check prompt logprobs
96
        # The first prompt logprob is always None, so we compare it from 1:.
97
98
99
        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():
100
                torch.testing.assert_close(logprob.logprob,
101
102
103
104
                                           hf_logprob[0][i][token_id].item(),
                                           atol=1e-2,
                                           rtol=1e-2)
        vllm_sample_logprobs = vllm_result.outputs[0].logprobs
105
106
107
        for i, top_logprobs in enumerate(vllm_sample_logprobs):
            for token_id, sample_logprob in top_logprobs.items():
                logprob = sample_logprob.logprob
108
109
110
111
                torch.testing.assert_close(logprob,
                                           hf_logprob[i][-1][token_id].item(),
                                           atol=1e-2,
                                           rtol=1e-2)
112
113
114
115
                if detokenize:
                    assert isinstance(sample_logprob.decoded_token, str), (
                        "The token should be decoded by the time it is returned"
                        " to the user.")
116

117
118
119
120
121
122
123
124
125
126
127
    # 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

128
129
130
131
132
133
134
135
136
137

def test_max_logprobs():
    runner = VllmRunner("facebook/opt-125m", max_logprobs=1)
    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)