conftest.py 18 KB
Newer Older
1
2
import contextlib
import gc
3
import os
4
from typing import Any, Dict, List, Optional, Tuple, TypeVar
Woosuk Kwon's avatar
Woosuk Kwon committed
5
6
7

import pytest
import torch
8
import torch.nn as nn
9
import torch.nn.functional as F
10
from PIL import Image
11
12
from transformers import (AutoModelForCausalLM, AutoModelForVision2Seq,
                          AutoProcessor, AutoTokenizer, BatchEncoding)
Woosuk Kwon's avatar
Woosuk Kwon committed
13
14

from vllm import LLM, SamplingParams
15
from vllm.config import TokenizerPoolConfig, VisionLanguageConfig
16
from vllm.distributed import destroy_model_parallel
17
from vllm.inputs import TextPrompt
18
from vllm.logger import init_logger
19
20
21
from vllm.multimodal import MultiModalData
from vllm.multimodal.image import ImageFeatureData, ImagePixelData
from vllm.sequence import SampleLogprobs
22
from vllm.utils import is_cpu
23
24

logger = init_logger(__name__)
Woosuk Kwon's avatar
Woosuk Kwon committed
25

26
27
28
_TEST_DIR = os.path.dirname(__file__)
_TEST_PROMPTS = [os.path.join(_TEST_DIR, "prompts", "example.txt")]
_LONG_PROMPTS = [os.path.join(_TEST_DIR, "prompts", "summary.txt")]
29

30
# Multi modal related
31
# You can use `.buildkite/download-images.sh` to download the assets
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
_PIXEL_VALUES_FILES = [
    os.path.join(_TEST_DIR, "images", filename) for filename in
    ["stop_sign_pixel_values.pt", "cherry_blossom_pixel_values.pt"]
]
_IMAGE_FEATURES_FILES = [
    os.path.join(_TEST_DIR, "images", filename) for filename in
    ["stop_sign_image_features.pt", "cherry_blossom_image_features.pt"]
]
_IMAGE_FILES = [
    os.path.join(_TEST_DIR, "images", filename)
    for filename in ["stop_sign.jpg", "cherry_blossom.jpg"]
]
_IMAGE_PROMPTS = [
    "<image>\nUSER: What's the content of the image?\nASSISTANT:",
    "<image>\nUSER: What is the season?\nASSISTANT:"
]
assert len(_PIXEL_VALUES_FILES) == len(_IMAGE_FEATURES_FILES) == len(
    _IMAGE_FILES) == len(_IMAGE_PROMPTS)

51

52
def _read_prompts(filename: str) -> List[str]:
53
    with open(filename, "r") as f:
54
55
        prompts = f.readlines()
        return prompts
Woosuk Kwon's avatar
Woosuk Kwon committed
56
57


58
59
60
61
62
def cleanup():
    destroy_model_parallel()
    with contextlib.suppress(AssertionError):
        torch.distributed.destroy_process_group()
    gc.collect()
63
64
    if not is_cpu():
        torch.cuda.empty_cache()
65
66


67
@pytest.fixture()
68
def should_do_global_cleanup_after_test(request) -> bool:
69
70
71
72
    """Allow subdirectories to skip global cleanup by overriding this fixture.
    This can provide a ~10x speedup for non-GPU unit tests since they don't need
    to initialize torch.
    """
73
74
75
76

    if request.node.get_closest_marker("skip_global_cleanup"):
        return False

77
78
79
    return True


80
@pytest.fixture(autouse=True)
81
def cleanup_fixture(should_do_global_cleanup_after_test: bool):
82
    yield
83
84
    if should_do_global_cleanup_after_test:
        cleanup()
85
86


87
88
89
90
91
92
93
94
95
96
97
@pytest.fixture(scope="session")
def hf_image_prompts() -> List[str]:
    return _IMAGE_PROMPTS


@pytest.fixture(scope="session")
def hf_images() -> List[Image.Image]:
    return [Image.open(filename) for filename in _IMAGE_FILES]


@pytest.fixture()
98
def vllm_images(request) -> List[MultiModalData]:
99
100
101
    vision_language_config = request.getfixturevalue("model_and_config")[1]
    if vision_language_config.image_input_type == (
            VisionLanguageConfig.ImageInputType.IMAGE_FEATURES):
102
103
104
105
        return [
            ImageFeatureData(torch.load(filename))
            for filename in _IMAGE_FEATURES_FILES
        ]
106
    else:
107
108
109
110
111
112
113
114
        return [
            ImagePixelData(Image.open(filename)) for filename in _IMAGE_FILES
        ]


@pytest.fixture()
def vllm_image_tensors(request) -> List[torch.Tensor]:
    return [torch.load(filename) for filename in _PIXEL_VALUES_FILES]
115
116
117
118
119
120
121
122
123
124
125


@pytest.fixture()
def vllm_image_prompts(request) -> List[str]:
    vision_language_config = request.getfixturevalue("model_and_config")[1]
    return [
        "<image>" * (vision_language_config.image_feature_size - 1) + p
        for p in _IMAGE_PROMPTS
    ]


Woosuk Kwon's avatar
Woosuk Kwon committed
126
127
@pytest.fixture
def example_prompts() -> List[str]:
128
129
    prompts = []
    for filename in _TEST_PROMPTS:
130
        prompts += _read_prompts(filename)
131
132
133
134
135
136
137
    return prompts


@pytest.fixture
def example_long_prompts() -> List[str]:
    prompts = []
    for filename in _LONG_PROMPTS:
138
        prompts += _read_prompts(filename)
139
    return prompts
Woosuk Kwon's avatar
Woosuk Kwon committed
140
141
142
143
144
145
146
147


_STR_DTYPE_TO_TORCH_DTYPE = {
    "half": torch.half,
    "bfloat16": torch.bfloat16,
    "float": torch.float,
}

148
_T = TypeVar("_T", nn.Module, torch.Tensor, BatchEncoding)
149

Woosuk Kwon's avatar
Woosuk Kwon committed
150
151
152

class HfRunner:

153
    def wrap_device(self, input: _T) -> _T:
154
155
156
157
158
        if not is_cpu():
            return input.to("cuda")
        else:
            return input.to("cpu")

Woosuk Kwon's avatar
Woosuk Kwon committed
159
160
161
162
    def __init__(
        self,
        model_name: str,
        dtype: str = "half",
163
164
165
        *,
        is_embedding_model: bool = False,
        is_vision_model: bool = False,
Woosuk Kwon's avatar
Woosuk Kwon committed
166
167
168
    ) -> None:
        assert dtype in _STR_DTYPE_TO_TORCH_DTYPE
        torch_dtype = _STR_DTYPE_TO_TORCH_DTYPE[dtype]
169

170
        self.model_name = model_name
171

172
        if is_embedding_model:
173
174
            # Lazy init required for AMD CI
            from sentence_transformers import SentenceTransformer
175
176
177
178
179
            self.model = self.wrap_device(
                SentenceTransformer(
                    model_name,
                    device="cpu",
                ).to(dtype=torch_dtype))
180
        else:
181
182
183
184
185
            if is_vision_model:
                auto_cls = AutoModelForVision2Seq
            else:
                auto_cls = AutoModelForCausalLM

186
            self.model = self.wrap_device(
187
                auto_cls.from_pretrained(
188
189
190
191
                    model_name,
                    torch_dtype=torch_dtype,
                    trust_remote_code=True,
                ))
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209

        self.tokenizer = AutoTokenizer.from_pretrained(
            model_name,
            torch_dtype=torch_dtype,
            trust_remote_code=True,
        )

        try:
            self.processor = AutoProcessor.from_pretrained(
                model_name,
                torch_dtype=torch_dtype,
                trust_remote_code=True,
            )
        except Exception:
            logger.warning(
                "Unable to auto-load processor from HuggingFace for "
                "model %s. Using tokenizer instead.", model_name)
            self.processor = self.tokenizer
Woosuk Kwon's avatar
Woosuk Kwon committed
210
211
212
213

    def generate(
        self,
        prompts: List[str],
214
        images: Optional[List[Image.Image]] = None,
Woosuk Kwon's avatar
Woosuk Kwon committed
215
        **kwargs,
216
    ) -> List[Tuple[List[List[int]], List[str]]]:
217
218
        if images:
            assert len(prompts) == len(images)
219
220

        outputs: List[Tuple[List[List[int]], List[str]]] = []
221
        for i, prompt in enumerate(prompts):
222
223
224
225
226
227
228
229
230
            processor_kwargs: Dict[str, Any] = {
                "text": prompt,
                "return_tensors": "pt",
            }
            if images is not None and images[i] is not None:
                processor_kwargs["images"] = images[i]

            inputs = self.processor(**processor_kwargs)

Woosuk Kwon's avatar
Woosuk Kwon committed
231
            output_ids = self.model.generate(
232
                **self.wrap_device(inputs),
Woosuk Kwon's avatar
Woosuk Kwon committed
233
234
235
                use_cache=True,
                **kwargs,
            )
236
            output_str = self.processor.batch_decode(
Woosuk Kwon's avatar
Woosuk Kwon committed
237
238
239
                output_ids,
                skip_special_tokens=True,
                clean_up_tokenization_spaces=False,
240
241
            )
            output_ids = output_ids.cpu().tolist()
Woosuk Kwon's avatar
Woosuk Kwon committed
242
243
244
245
246
247
248
            outputs.append((output_ids, output_str))
        return outputs

    def generate_greedy(
        self,
        prompts: List[str],
        max_tokens: int,
249
        images: Optional[List[Image.Image]] = None,
Woosuk Kwon's avatar
Woosuk Kwon committed
250
    ) -> List[Tuple[List[int], str]]:
251
252
        outputs = self.generate(prompts,
                                do_sample=False,
253
254
                                max_new_tokens=max_tokens,
                                images=images)
255
256
257

        return [(output_ids[0], output_str[0])
                for output_ids, output_str in outputs]
258
259
260
261
262
263

    def generate_beam_search(
        self,
        prompts: List[str],
        beam_width: int,
        max_tokens: int,
264
    ) -> List[Tuple[List[List[int]], List[str]]]:
265
266
267
268
269
270
271
272
273
274
275
276
277
278
        outputs = self.generate(prompts,
                                do_sample=False,
                                max_new_tokens=max_tokens,
                                num_beams=beam_width,
                                num_return_sequences=beam_width)
        for i in range(len(outputs)):
            output_ids, output_str = outputs[i]
            for j in range(len(output_ids)):
                output_ids[j] = [
                    x for x in output_ids[j]
                    if x != self.tokenizer.pad_token_id
                ]
            outputs[i] = (output_ids, output_str)
        return outputs
Woosuk Kwon's avatar
Woosuk Kwon committed
279

280
281
282
283
284
285
286
287
288
    def generate_greedy_logprobs(
        self,
        prompts: List[str],
        max_tokens: int,
    ) -> List[List[torch.Tensor]]:
        all_logprobs = []
        for prompt in prompts:
            input_ids = self.tokenizer(prompt, return_tensors="pt").input_ids
            output = self.model.generate(
289
                self.wrap_device(input_ids),
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
                use_cache=True,
                do_sample=False,
                max_new_tokens=max_tokens,
                output_hidden_states=True,
                return_dict_in_generate=True,
            )
            seq_logprobs = []
            for hidden_states in output.hidden_states:
                last_hidden_states = hidden_states[-1][0]
                logits = torch.matmul(
                    last_hidden_states,
                    self.model.get_output_embeddings().weight.t(),
                )
                if self.model.get_output_embeddings().bias is not None:
                    logits += self.model.get_output_embeddings(
                    ).bias.unsqueeze(0)
306
                logprobs = F.log_softmax(logits, dim=-1, dtype=torch.float32)
307
308
309
310
                seq_logprobs.append(logprobs)
            all_logprobs.append(seq_logprobs)
        return all_logprobs

311
312
313
314
315
    def generate_greedy_logprobs_limit(
        self,
        prompts: List[str],
        max_tokens: int,
        num_logprobs: int,
316
317
318
319
    ) -> List[Tuple[List[int], str, List[Dict[int, float]]]]:
        all_logprobs: List[List[Dict[int, float]]] = []
        all_output_ids: List[List[int]] = []
        all_output_strs: List[str] = []
320
321
322
323

        for prompt in prompts:
            input_ids = self.tokenizer(prompt, return_tensors="pt").input_ids
            output = self.model.generate(
324
                self.wrap_device(input_ids),
325
326
327
328
329
330
331
                use_cache=True,
                do_sample=False,
                max_new_tokens=max_tokens,
                output_hidden_states=True,
                return_dict_in_generate=True,
            )

332
            seq_logprobs: List[torch.Tensor] = []
333
334
335
336
337
338
339
340
341
342
            for _, hidden_states in enumerate(output.hidden_states):
                last_hidden_states = hidden_states[-1][0]
                logits = torch.matmul(
                    last_hidden_states,
                    self.model.get_output_embeddings().weight.t(),
                )
                if getattr(self.model.get_output_embeddings(), "bias",
                           None) is not None:
                    logits += self.model.get_output_embeddings(
                    ).bias.unsqueeze(0)
343
                logprobs = F.log_softmax(logits, dim=-1, dtype=torch.float32)
344
345
346
                seq_logprobs.append(logprobs)

            # convert to dict
347
            seq_logprobs_lst: List[Dict[int, float]] = []
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
            for tok_idx, tok_logprobs in enumerate(seq_logprobs):
                # drop prompt logprobs
                if tok_idx == 0:
                    tok_logprobs = tok_logprobs[-1, :].reshape(1, -1)
                topk = tok_logprobs.topk(num_logprobs)

                tok_logprobs_dct = {}
                for token_id, logprob in zip(topk.indices[0], topk.values[0]):
                    tok_logprobs_dct[token_id.item()] = logprob.item()

                seq_logprobs_lst.append(tok_logprobs_dct)

            all_logprobs.append(seq_logprobs_lst)
            seq_ids = output.sequences[0]
            output_len = seq_ids.shape[0] - input_ids.shape[1]
            output_ids = seq_ids[-output_len:]
            all_output_ids.append(output_ids.tolist())
            all_output_strs.append(self.tokenizer.decode(output_ids))

        outputs = zip(all_output_ids, all_output_strs, all_logprobs)
        return [(output_ids, output_str, output_logprobs)
                for output_ids, output_str, output_logprobs in outputs]

371
372
373
    def encode(self, prompts: List[str]) -> List[List[torch.Tensor]]:
        return self.model.encode(prompts)

374
375
376
377
    def __del__(self):
        del self.model
        cleanup()

Woosuk Kwon's avatar
Woosuk Kwon committed
378
379
380
381
382
383
384
385
386
387
388
389

@pytest.fixture
def hf_runner():
    return HfRunner


class VllmRunner:

    def __init__(
        self,
        model_name: str,
        tokenizer_name: Optional[str] = None,
390
391
        # Use smaller max model length, otherwise bigger model cannot run due
        # to kv cache size limit.
392
        max_model_len: int = 1024,
Woosuk Kwon's avatar
Woosuk Kwon committed
393
        dtype: str = "half",
394
        disable_log_stats: bool = True,
395
        tensor_parallel_size: int = 1,
396
397
        block_size: int = 16,
        enable_chunked_prefill: bool = False,
398
        swap_space: int = 4,
399
        **kwargs,
Woosuk Kwon's avatar
Woosuk Kwon committed
400
401
402
403
404
405
    ) -> None:
        self.model = LLM(
            model=model_name,
            tokenizer=tokenizer_name,
            trust_remote_code=True,
            dtype=dtype,
406
            swap_space=swap_space,
407
            disable_log_stats=disable_log_stats,
408
            tensor_parallel_size=tensor_parallel_size,
409
            max_model_len=max_model_len,
410
411
            block_size=block_size,
            enable_chunked_prefill=enable_chunked_prefill,
412
            **kwargs,
Woosuk Kwon's avatar
Woosuk Kwon committed
413
414
415
416
417
418
        )

    def generate(
        self,
        prompts: List[str],
        sampling_params: SamplingParams,
419
        images: Optional[List[MultiModalData]] = None,
420
    ) -> List[Tuple[List[List[int]], List[str]]]:
421
        if images is not None:
422
            assert len(prompts) == len(images)
423

424
425
426
427
        inputs = [TextPrompt(prompt=prompt) for prompt in prompts]
        if images is not None:
            for i, image in enumerate(images):
                inputs[i]["multi_modal_data"] = image
428

429
        req_outputs = self.model.generate(inputs,
430
                                          sampling_params=sampling_params)
431
432

        outputs: List[Tuple[List[List[int]], List[str]]] = []
Woosuk Kwon's avatar
Woosuk Kwon committed
433
434
435
        for req_output in req_outputs:
            prompt_str = req_output.prompt
            prompt_ids = req_output.prompt_token_ids
436
437
            req_sample_output_ids: List[List[int]] = []
            req_sample_output_strs: List[str] = []
438
439
440
441
442
443
            for sample in req_output.outputs:
                output_str = sample.text
                output_ids = sample.token_ids
                req_sample_output_ids.append(prompt_ids + output_ids)
                req_sample_output_strs.append(prompt_str + output_str)
            outputs.append((req_sample_output_ids, req_sample_output_strs))
Woosuk Kwon's avatar
Woosuk Kwon committed
444
445
        return outputs

446
447
448
449
    def generate_w_logprobs(
        self,
        prompts: List[str],
        sampling_params: SamplingParams,
450
    ) -> List[Tuple[List[int], str, Optional[SampleLogprobs]]]:
451
452
453
454
        assert sampling_params.logprobs is not None

        req_outputs = self.model.generate(prompts,
                                          sampling_params=sampling_params)
455
        outputs: List[Tuple[List[int], str, Optional[SampleLogprobs]]] = []
456
457
458
459
460
461
462
463
        for req_output in req_outputs:
            for sample in req_output.outputs:
                output_str = sample.text
                output_ids = sample.token_ids
                output_logprobs = sample.logprobs
            outputs.append((output_ids, output_str, output_logprobs))
        return outputs

Woosuk Kwon's avatar
Woosuk Kwon committed
464
465
466
467
    def generate_greedy(
        self,
        prompts: List[str],
        max_tokens: int,
468
        images: Optional[List[MultiModalData]] = None,
Woosuk Kwon's avatar
Woosuk Kwon committed
469
470
    ) -> List[Tuple[List[int], str]]:
        greedy_params = SamplingParams(temperature=0.0, max_tokens=max_tokens)
471
        outputs = self.generate(prompts, greedy_params, images=images)
472
473
        return [(output_ids[0], output_str[0])
                for output_ids, output_str in outputs]
474

475
476
477
478
479
    def generate_greedy_logprobs(
        self,
        prompts: List[str],
        max_tokens: int,
        num_logprobs: int,
480
    ) -> List[Tuple[List[int], str, Optional[SampleLogprobs]]]:
481
482
483
484
485
486
487
488
        greedy_logprobs_params = SamplingParams(temperature=0.0,
                                                max_tokens=max_tokens,
                                                logprobs=num_logprobs)
        outputs = self.generate_w_logprobs(prompts, greedy_logprobs_params)

        return [(output_ids, output_str, output_logprobs)
                for output_ids, output_str, output_logprobs in outputs]

489
490
491
492
493
    def generate_beam_search(
        self,
        prompts: List[str],
        beam_width: int,
        max_tokens: int,
494
    ) -> List[Tuple[List[List[int]], List[str]]]:
495
496
497
498
499
500
        beam_search_params = SamplingParams(n=beam_width,
                                            use_beam_search=True,
                                            temperature=0.0,
                                            max_tokens=max_tokens)
        outputs = self.generate(prompts, beam_search_params)
        return outputs
Woosuk Kwon's avatar
Woosuk Kwon committed
501

502
503
504
505
506
507
508
509
    def encode(self, prompts: List[str]) -> List[List[float]]:
        req_outputs = self.model.encode(prompts)
        outputs = []
        for req_output in req_outputs:
            embedding = req_output.outputs.embedding
            outputs.append(embedding)
        return outputs

510
511
512
513
    def __del__(self):
        del self.model
        cleanup()

Woosuk Kwon's avatar
Woosuk Kwon committed
514

515
@pytest.fixture(scope="session")
Woosuk Kwon's avatar
Woosuk Kwon committed
516
517
def vllm_runner():
    return VllmRunner
518
519
520
521
522
523
524
525
526
527


def get_tokenizer_pool_config(tokenizer_group_type):
    if tokenizer_group_type is None:
        return None
    if tokenizer_group_type == "ray":
        return TokenizerPoolConfig(pool_size=1,
                                   pool_type="ray",
                                   extra_config={})
    raise ValueError(f"Unknown tokenizer_group_type: {tokenizer_group_type}")
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543


@pytest.fixture()
def temporary_enable_log_propagate():
    import logging
    logger = logging.getLogger("vllm")
    logger.propagate = True
    yield
    logger.propagate = False


@pytest.fixture()
def caplog_vllm(temporary_enable_log_propagate, caplog):
    # To capture vllm log, we should enable propagate=True temporarily
    # because caplog depends on logs propagated to the root logger.
    yield caplog