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

4
import warnings
5
from collections.abc import Sequence
6
from typing import Any, NamedTuple, Optional, Union
7

8
import torch
9
import torch.nn.functional as F
10

11
from vllm.config import ModelConfig, RunnerOption
12
from vllm.inputs import InputContext
13
from vllm.sequence import Logprob, PromptLogprobs, SampleLogprobs
14

15
16
from .registry import HF_EXAMPLE_MODELS

17
TokensText = tuple[list[int], str]
18
19


20
21
22
23
24
25
26
def check_outputs_equal(
    *,
    outputs_0_lst: Sequence[TokensText],
    outputs_1_lst: Sequence[TokensText],
    name_0: str,
    name_1: str,
):
27
    """
28
    Compare the two sequences generated by different models,
29
30
31
32
33
34
35
36
37
38
    which should be equal.
    """
    assert len(outputs_0_lst) == len(outputs_1_lst)

    for prompt_idx, (outputs_0,
                     outputs_1) in enumerate(zip(outputs_0_lst,
                                                 outputs_1_lst)):
        output_ids_0, output_str_0 = outputs_0
        output_ids_1, output_str_1 = outputs_1

39
40
41
42
43
44
45
        # The text and token outputs should exactly match
        fail_msg = (f"Test{prompt_idx}:"
                    f"\n{name_0}:\t{output_str_0!r}"
                    f"\n{name_1}:\t{output_str_1!r}")

        assert output_str_0 == output_str_1, fail_msg
        assert output_ids_0 == output_ids_1, fail_msg
46
47


48
49
50
51
52
53
# Representation of generated sequence as a tuple of
# * Token ID list
# * String
# * List of top sample logprobs for each sampled token
#
# Assumes prompt logprobs were not requested.
54
TokensTextLogprobs = tuple[list[int], str, Optional[Union[list[dict[int,
55
56
                                                                    float]],
                                                          SampleLogprobs]]]
57

58
59
60
61
62
63
64
# Allow for tokens to be represented as str's rather than IDs;
# tuple of
# * Token string representations list
# * String
# * Optional list of top sample logprobs for each sampled token
#
# Assumes prompt logprobs were not requested.
65
66
TextTextLogprobs = tuple[list[str], str, Optional[Union[list[dict[str, float]],
                                                        list[dict[str,
67
68
                                                                  Logprob]]]]]

69
70
71
72
73
74
75
# Representation of generated sequence as a tuple of
# * Token ID list
# * String
# * Optional list of top sample logprobs for each sampled token
# * Optional list of top prompt logprobs for each prompt token
#
# Allows prompt logprobs to be requested.
76
77
78
TokensTextLogprobsPromptLogprobs = tuple[
    list[int], str, Optional[Union[list[dict[int, float]], SampleLogprobs]],
    Optional[Union[list[Optional[dict[int, float]]], PromptLogprobs]]]
79

80

81
82
def check_logprobs_close(
    *,
83
84
85
86
87
88
    outputs_0_lst: Sequence[Union[TokensTextLogprobs,
                                  TokensTextLogprobsPromptLogprobs,
                                  TextTextLogprobs]],
    outputs_1_lst: Sequence[Union[TokensTextLogprobs,
                                  TokensTextLogprobsPromptLogprobs,
                                  TextTextLogprobs]],
89
90
    name_0: str,
    name_1: str,
91
    num_outputs_0_skip_tokens: int = 0,
92
    warn_on_mismatch: bool = True,
93
94
95
    always_check_logprobs: bool = False,
) -> None:
    """Compare the logprobs of two sequences generated by different models,
96
    which should be similar but not necessarily equal.
97

98
99
100
101
102
103
104
105
106
107
108
109
    How sample logprobs are compared:
    * `always_check_logprobs == True`: set of highest-logprob token ids
      must match between seq0 and seq1 at all sampled token offsets
    * `always_check_logprobs == False`: highest-logprob token ids are
      only compared at sampled token offsets for which generated token
      ids don't match

    Prompt logprobs must be provided either for both input sequences, or
    for neither. If prompt logprobs are provided, then highest-logprob
    prompt token ids must match between seq0 and seq1 at all prompt token
    offsets.

110
111
112
113
114
115
    Args:
      outputs_0_lst: First sequence to compare
      outputs_0_lst: Second sequence to compare
      name_0: sequence #0 name
      name_1: sequence #1 name
      num_outputs_0_skip_tokens: If > 0, specifies the number of initial
116
117
118
119
120
                                 sequence #0 tokens & logprobs to discard
                                 before comparison, i.e. all
                                 of sequence #1 will be compared to
                                 sequence #0 beginning at index
                                 num_outputs_0_skip_tokens
121
      warn_on_mismatch: Issue a warning if there is token-wise or text-wise
122
                        mismatch between the two sequences
123
      always_check_logprobs: If true, check logprobs even when tokens match
124
    """
125
126
    assert len(outputs_0_lst) == len(outputs_1_lst)

127
128
129
130
    # Loop through responses to each prompt.
    for prompt_idx, (outputs_0,
                     outputs_1) in enumerate(zip(outputs_0_lst,
                                                 outputs_1_lst)):
131
132
133
134
135
136
137
138
139
140
141
142
143
144
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
177
178
179
180
181
182
183
184
185
186
187
188
189
        assert len(outputs_0) == len(outputs_1)
        if len(outputs_0) == 3:
            assert len(outputs_1) == 3
            # Break out tokens, text & sample logprobs
            # (prompt logprobs were not provided)
            output_ids_0, output_str_0, logprobs_0 = outputs_0
            output_ids_1, output_str_1, logprobs_1 = outputs_1
        elif len(outputs_0) == 4:
            assert len(outputs_1) == 4
            # Break out tokens, text, sample logprobs & prompt logprobs
            (
                output_ids_0,
                output_str_0,
                logprobs_0,
                prompt_logprobs_0,
            ) = outputs_0
            (
                output_ids_1,
                output_str_1,
                logprobs_1,
                prompt_logprobs_1,
            ) = outputs_1

            # Test prompt logprobs closeness
            if (prompt_logprobs_0 is not None
                    and prompt_logprobs_1 is not None):
                # Both sequences' prompt logprobs lists are not `None``
                # (although individual list elements may be `None`);
                # for each token's logprobs:
                for idx, (logprobs_elem_0, logprobs_elem_1) in enumerate(
                        zip(prompt_logprobs_0, prompt_logprobs_1)):
                    fail_msg = (
                        f"Prompt logprobs test:"
                        f"\n{name_0}:\tPrompt index {idx}\t{logprobs_elem_0}"
                        f"\n{name_1}:\tPrompt index {idx}\t{logprobs_elem_1}")

                    if logprobs_elem_0 is None:
                        # If the seq 0 token's logprobs are `None`,
                        # the seq 1 token's logprobs must be `None`
                        assert logprobs_elem_1 is None, fail_msg
                    else:
                        # If the seq 0 token's logprobs are not `None`,
                        # the seq 1 token's logprobs must not be `None`
                        assert logprobs_elem_1 is not None, fail_msg
                        # Logprobs check: top-k token choices must be the same
                        assert (set(logprobs_elem_0.keys()) == set(
                            logprobs_elem_1.keys())), fail_msg
            else:
                # Both sequence logprobs lists must be `None`
                fail_msg = (f"Prompt logprobs test:"
                            f"\n{name_0}:\tlogprobs\t{prompt_logprobs_0}"
                            f"\n{name_1}:\tlogprobs\t{prompt_logprobs_1}")

                assert (prompt_logprobs_0 is None
                        and prompt_logprobs_1 is None), fail_msg
        else:
            raise ValueError(f"Outputs tuple must have 3 or 4 elements but "
                             f"{len(outputs_0)} elements were provided: "
                             f"{outputs_0}")
190

191
192
193
194
195
        if logprobs_0 is None:
            logprobs_0 = [None] * len(output_ids_0)
        if logprobs_1 is None:
            logprobs_1 = [None] * len(output_ids_1)

196
197
198
199
200
201
202
203
204
        # Skip specified number of initial sequence #0 tokens
        # & logprobs, leaving output text as-is for simplicity
        # (text mismatches may generate warnings but do not
        # cause the test to fail.)
        if num_outputs_0_skip_tokens < 0:
            raise ValueError("num_outputs_0_skip_tokens must be non-negative")
        output_ids_0 = output_ids_0[num_outputs_0_skip_tokens:]
        logprobs_0 = logprobs_0[num_outputs_0_skip_tokens:]

205
206
207
208
        # Loop through generated tokens.
        for idx, (output_id_0,
                  output_id_1) in enumerate(zip(output_ids_0, output_ids_1)):

209
210
211
212
213
214
            is_tok_mismatch = output_id_0 != output_id_1

            # If generated tokens don't match
            # or it is desired to always check logprobs,
            # then
            if is_tok_mismatch or always_check_logprobs:
215
216
217
                logprobs_elem_0 = logprobs_0[idx]
                logprobs_elem_1 = logprobs_1[idx]

218
                # Each predicted token must be in top N logprobs of the other
219
                fail_msg = (
220
                    f"Test{prompt_idx}:"
221
                    f"\nMatched tokens:\t{output_ids_0[:idx]}"
222
223
224
225
226
227
228
229
                    f"\n{name_0}:\t{output_str_0!r}\t{logprobs_elem_0}"
                    f"\n{name_1}:\t{output_str_1!r}\t{logprobs_elem_1}")

                assert logprobs_elem_0 is not None, fail_msg
                assert logprobs_elem_1 is not None, fail_msg
                assert output_id_0 in logprobs_elem_1, fail_msg
                assert output_id_1 in logprobs_elem_0, fail_msg

230
                if warn_on_mismatch and is_tok_mismatch:
231
232
233
234
235
236
                    with warnings.catch_warnings():
                        # This ensures that repeated warnings are shown
                        # in the output, not just the first occurrence
                        warnings.simplefilter("always")

                        warnings.warn(fail_msg, stacklevel=2)
237
238
239

                # Break out since sequences will now diverge.
                break
240
241
242
243
244
245
246
247
248
249
250
251
252
253
        else:
            if output_str_0 != output_str_1 and warn_on_mismatch:
                # The token outputs exactly match,
                # so the text outputs should exactly match as well
                fail_msg = (f"Test{prompt_idx}:"
                            f"\n{name_0}:\t{output_str_0!r}"
                            f"\n{name_1}:\t{output_str_1!r}")

                with warnings.catch_warnings():
                    # This ensures that repeated warnings are shown
                    # in the output, not just the first occurrence
                    warnings.simplefilter("always")

                    warnings.warn(fail_msg, stacklevel=2)
254
255


256
def build_model_context(
257
    model_id: str,
258
    runner: RunnerOption = "auto",
259
    dtype: Union[str, torch.dtype] = "auto",
260
    model_config_kwargs: Optional[dict[str, Any]] = None,
261
262
    mm_processor_kwargs: Optional[dict[str, Any]] = None,
    limit_mm_per_prompt: Optional[dict[str, int]] = None,
263
264
    disable_mm_preprocessor_cache: bool = True,
):
265
    """Creates an InputContext for a given model.
266

267
    Args:
268
        model_id: ID of the model being considered.
269
270
271
272
273
274
275
        mm_processor_kwargs: optional processor kwargs for to be leveraged
            in the input processor, mapper, dummy data creation, etc.
        limit_mm_per_prompt: Multimodal limits.

    Returns:
        InputContext for the model being considered.
    """
276
277
278
279
    model_info = HF_EXAMPLE_MODELS.find_hf_info(model_id)
    model_info.check_available_online(on_fail="skip")
    model_info.check_transformers_version(on_fail="skip")

280
    model_config_kwargs = model_config_kwargs or {}
281
    model_config = ModelConfig(
282
        model_id,
283
        runner=runner,
284
285
        tokenizer=model_info.tokenizer or model_id,
        tokenizer_mode=model_info.tokenizer_mode,
286
        revision=model_info.revision,
287
        trust_remote_code=model_info.trust_remote_code,
288
        dtype=dtype,
289
290
291
        seed=0,
        mm_processor_kwargs=mm_processor_kwargs,
        limit_mm_per_prompt=limit_mm_per_prompt,
292
        disable_mm_preprocessor_cache=disable_mm_preprocessor_cache,
293
        hf_overrides=model_info.hf_overrides,
294
        **model_config_kwargs,
295
296
    )
    return InputContext(model_config)
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318


def check_embeddings_close(
    *,
    embeddings_0_lst: Sequence[list[float]],
    embeddings_1_lst: Sequence[list[float]],
    name_0: str,
    name_1: str,
    tol: float = 1e-3,
) -> None:
    assert len(embeddings_0_lst) == len(embeddings_1_lst)

    for prompt_idx, (embeddings_0, embeddings_1) in enumerate(
            zip(embeddings_0_lst, embeddings_1_lst)):
        assert len(embeddings_0) == len(embeddings_1), (
            f"Length mismatch: {len(embeddings_0)} vs. {len(embeddings_1)}")

        sim = F.cosine_similarity(torch.tensor(embeddings_0),
                                  torch.tensor(embeddings_1),
                                  dim=0)

        fail_msg = (f"Test{prompt_idx}:"
319
                    f"\nCosine similarity: \t{sim:.4f}"
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
                    f"\n{name_0}:\t{embeddings_0[:16]!r}"
                    f"\n{name_1}:\t{embeddings_1[:16]!r}")

        assert sim >= 1 - tol, fail_msg


def matryoshka_fy(tensor: torch.Tensor, dimensions: int):
    tensor = torch.tensor(tensor)
    tensor = tensor[..., :dimensions]
    tensor = F.normalize(tensor, p=2, dim=1)
    return tensor


class EmbedModelInfo(NamedTuple):
    name: str
335
    is_matryoshka: bool = False
336
337
    matryoshka_dimensions: Optional[list[int]] = None
    architecture: str = ""
338
    dtype: str = "auto"
339
    enable_test: bool = True
340
341
342
343
344
345
346


class RerankModelInfo(NamedTuple):
    name: str
    architecture: str = ""
    dtype: str = "auto"
    enable_test: bool = True