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

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

from vllm import LLM, SamplingParams
17
from vllm.assets.image import ImageAsset
18
from vllm.config import TokenizerPoolConfig
19
20
from vllm.distributed import (destroy_distributed_environment,
                              destroy_model_parallel)
21
from vllm.inputs import TextPrompt
22
from vllm.logger import init_logger
23
24
from vllm.sequence import SampleLogprobs
from vllm.utils import cuda_device_count_stateless, is_cpu
25

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

28
29
30
_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")]
31
32


33
def _read_prompts(filename: str) -> List[str]:
34
    with open(filename, "r") as f:
35
36
        prompts = f.readlines()
        return prompts
Woosuk Kwon's avatar
Woosuk Kwon committed
37
38


39
40
41
class _ImageAssetPrompts(TypedDict):
    stop_sign: str
    cherry_blossom: str
42
43
44
45
46
47
48


if sys.version_info < (3, 9):
    # UserList cannot be subscripted
    class _ImageAssetsBase(UserList):
        pass
else:
49

50
51
    class _ImageAssetsBase(UserList[ImageAsset]):
        pass
52

53
54

class _ImageAssets(_ImageAssetsBase):
55
56

    def __init__(self) -> None:
57
58
59
60
        super().__init__([
            ImageAsset("stop_sign"),
            ImageAsset("cherry_blossom"),
        ])
61
62
63
64
65
66
67
68

    def prompts(self, prompts: _ImageAssetPrompts) -> List[str]:
        """
        Convenience method to define the prompt for each test image.

        The order of the returned prompts matches the order of the
        assets when iterating through this object.
        """
69
        return [prompts["stop_sign"], prompts["cherry_blossom"]]
70
71
72
73
74
75


IMAGE_ASSETS = _ImageAssets()
"""Singleton instance of :class:`_ImageAssets`."""


76
77
def cleanup():
    destroy_model_parallel()
78
    destroy_distributed_environment()
79
80
81
    with contextlib.suppress(AssertionError):
        torch.distributed.destroy_process_group()
    gc.collect()
82
83
    if not is_cpu():
        torch.cuda.empty_cache()
84
85


86
@pytest.fixture()
87
def should_do_global_cleanup_after_test(request) -> bool:
88
89
90
91
    """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.
    """
92
93
94
95

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

96
97
98
    return True


99
@pytest.fixture(autouse=True)
100
def cleanup_fixture(should_do_global_cleanup_after_test: bool):
101
    yield
102
103
    if should_do_global_cleanup_after_test:
        cleanup()
104
105


Woosuk Kwon's avatar
Woosuk Kwon committed
106
107
@pytest.fixture
def example_prompts() -> List[str]:
108
109
    prompts = []
    for filename in _TEST_PROMPTS:
110
        prompts += _read_prompts(filename)
111
112
113
114
115
116
117
    return prompts


@pytest.fixture
def example_long_prompts() -> List[str]:
    prompts = []
    for filename in _LONG_PROMPTS:
118
        prompts += _read_prompts(filename)
119
    return prompts
Woosuk Kwon's avatar
Woosuk Kwon committed
120
121


122
123
124
125
126
@pytest.fixture(scope="session")
def image_assets() -> _ImageAssets:
    return IMAGE_ASSETS


Woosuk Kwon's avatar
Woosuk Kwon committed
127
128
129
130
131
132
_STR_DTYPE_TO_TORCH_DTYPE = {
    "half": torch.half,
    "bfloat16": torch.bfloat16,
    "float": torch.float,
}

133
_T = TypeVar("_T", nn.Module, torch.Tensor, BatchEncoding)
134

Woosuk Kwon's avatar
Woosuk Kwon committed
135
136
137

class HfRunner:

138
    def wrap_device(self, input: _T) -> _T:
139
140
141
142
143
        if not is_cpu():
            return input.to("cuda")
        else:
            return input.to("cpu")

Woosuk Kwon's avatar
Woosuk Kwon committed
144
145
146
147
    def __init__(
        self,
        model_name: str,
        dtype: str = "half",
148
        *,
149
        model_kwargs: Optional[Dict[str, Any]] = None,
150
151
        is_embedding_model: bool = False,
        is_vision_model: bool = False,
152
        is_sparseml_model: bool = False,
Woosuk Kwon's avatar
Woosuk Kwon committed
153
154
155
    ) -> None:
        assert dtype in _STR_DTYPE_TO_TORCH_DTYPE
        torch_dtype = _STR_DTYPE_TO_TORCH_DTYPE[dtype]
156

157
        self.model_name = model_name
158

159
        if is_embedding_model:
160
161
            # Lazy init required for AMD CI
            from sentence_transformers import SentenceTransformer
162
163
164
165
166
            self.model = self.wrap_device(
                SentenceTransformer(
                    model_name,
                    device="cpu",
                ).to(dtype=torch_dtype))
167
        else:
168
169
            if is_vision_model:
                auto_cls = AutoModelForVision2Seq
170
171
172
            elif is_sparseml_model:
                from sparseml.transformers import SparseAutoModelForCausalLM
                auto_cls = SparseAutoModelForCausalLM
173
174
175
            else:
                auto_cls = AutoModelForCausalLM

176
            model_kwargs = model_kwargs if model_kwargs is not None else {}
177
            self.model = self.wrap_device(
178
                auto_cls.from_pretrained(
179
180
181
                    model_name,
                    torch_dtype=torch_dtype,
                    trust_remote_code=True,
182
                    **model_kwargs,
183
                ))
184
185
186
187
188
189
190
191

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

        try:
192
193
194
            # don't put this import at the top level
            # it will call torch.cuda.device_count()
            from transformers import AutoProcessor  # noqa: F401
195
196
197
198
199
200
201
202
203
204
            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
205
206
207
208

    def generate(
        self,
        prompts: List[str],
209
        images: Optional[List[Image.Image]] = None,
210
        **kwargs: Any,
211
    ) -> List[Tuple[List[List[int]], List[str]]]:
212
213
        if images:
            assert len(prompts) == len(images)
214
215

        outputs: List[Tuple[List[List[int]], List[str]]] = []
216
        for i, prompt in enumerate(prompts):
217
218
219
220
221
222
223
224
225
            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
226
            output_ids = self.model.generate(
227
                **self.wrap_device(inputs),
Woosuk Kwon's avatar
Woosuk Kwon committed
228
229
230
                use_cache=True,
                **kwargs,
            )
231
            output_str = self.processor.batch_decode(
Woosuk Kwon's avatar
Woosuk Kwon committed
232
233
234
                output_ids,
                skip_special_tokens=True,
                clean_up_tokenization_spaces=False,
235
236
            )
            output_ids = output_ids.cpu().tolist()
Woosuk Kwon's avatar
Woosuk Kwon committed
237
238
239
240
241
242
243
            outputs.append((output_ids, output_str))
        return outputs

    def generate_greedy(
        self,
        prompts: List[str],
        max_tokens: int,
244
        images: Optional[List[Image.Image]] = None,
245
        **kwargs: Any,
Woosuk Kwon's avatar
Woosuk Kwon committed
246
    ) -> List[Tuple[List[int], str]]:
247
248
        outputs = self.generate(prompts,
                                do_sample=False,
249
                                max_new_tokens=max_tokens,
Chang Su's avatar
Chang Su committed
250
251
                                images=images,
                                **kwargs)
252
253
254

        return [(output_ids[0], output_str[0])
                for output_ids, output_str in outputs]
255
256
257
258
259
260

    def generate_beam_search(
        self,
        prompts: List[str],
        beam_width: int,
        max_tokens: int,
261
    ) -> List[Tuple[List[List[int]], List[str]]]:
262
263
264
265
266
267
268
269
270
271
272
273
274
275
        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
276

277
278
279
280
    def generate_greedy_logprobs(
        self,
        prompts: List[str],
        max_tokens: int,
281
282
        images: Optional[List[Image.Image]] = None,
        **kwargs: Any,
283
    ) -> List[List[torch.Tensor]]:
284
285
286
287
288
289
290
291
292
293
294
        all_logprobs: List[List[torch.Tensor]] = []
        for i, prompt in enumerate(prompts):
            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)

295
            output = self.model.generate(
296
                **self.wrap_device(inputs),
297
298
299
300
301
                use_cache=True,
                do_sample=False,
                max_new_tokens=max_tokens,
                output_hidden_states=True,
                return_dict_in_generate=True,
302
                **kwargs,
303
            )
304
            seq_logprobs: List[torch.Tensor] = []
305
306
307
308
309
310
311
312
313
            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)
314
                logprobs = F.log_softmax(logits, dim=-1, dtype=torch.float32)
315
316
317
318
                seq_logprobs.append(logprobs)
            all_logprobs.append(seq_logprobs)
        return all_logprobs

319
320
321
322
323
    def generate_greedy_logprobs_limit(
        self,
        prompts: List[str],
        max_tokens: int,
        num_logprobs: int,
324
325
        images: Optional[List[Image.Image]] = None,
        **kwargs: Any,
326
327
328
329
    ) -> 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] = []
330

331
332
333
334
335
336
337
338
339
340
341
        for i, prompt in enumerate(prompts):
            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)
            input_ids = inputs.input_ids

342
            output = self.model.generate(
343
                **self.wrap_device(inputs),
344
345
346
347
348
                use_cache=True,
                do_sample=False,
                max_new_tokens=max_tokens,
                output_hidden_states=True,
                return_dict_in_generate=True,
349
                **kwargs,
350
351
            )

352
            seq_logprobs: List[torch.Tensor] = []
353
354
355
356
357
358
359
360
361
362
            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)
363
                logprobs = F.log_softmax(logits, dim=-1, dtype=torch.float32)
364
365
366
                seq_logprobs.append(logprobs)

            # convert to dict
367
            seq_logprobs_lst: List[Dict[int, float]] = []
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
            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]

391
392
393
    def encode(self, prompts: List[str]) -> List[List[torch.Tensor]]:
        return self.model.encode(prompts)

394
395
396
397
    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
398
399
400
        del self.model
        cleanup()

Woosuk Kwon's avatar
Woosuk Kwon committed
401

Cyrus Leung's avatar
Cyrus Leung committed
402
@pytest.fixture(scope="session")
Woosuk Kwon's avatar
Woosuk Kwon committed
403
404
405
406
407
408
409
410
411
412
def hf_runner():
    return HfRunner


class VllmRunner:

    def __init__(
        self,
        model_name: str,
        tokenizer_name: Optional[str] = None,
413
414
        # Use smaller max model length, otherwise bigger model cannot run due
        # to kv cache size limit.
415
        max_model_len: int = 1024,
Woosuk Kwon's avatar
Woosuk Kwon committed
416
        dtype: str = "half",
417
        disable_log_stats: bool = True,
418
        tensor_parallel_size: int = 1,
419
420
        block_size: int = 16,
        enable_chunked_prefill: bool = False,
421
        swap_space: int = 4,
Cyrus Leung's avatar
Cyrus Leung committed
422
        enforce_eager: bool = False,
423
        **kwargs,
Woosuk Kwon's avatar
Woosuk Kwon committed
424
425
426
427
428
429
    ) -> None:
        self.model = LLM(
            model=model_name,
            tokenizer=tokenizer_name,
            trust_remote_code=True,
            dtype=dtype,
430
            swap_space=swap_space,
Cyrus Leung's avatar
Cyrus Leung committed
431
            enforce_eager=enforce_eager,
432
            disable_log_stats=disable_log_stats,
433
            tensor_parallel_size=tensor_parallel_size,
434
            max_model_len=max_model_len,
435
436
            block_size=block_size,
            enable_chunked_prefill=enable_chunked_prefill,
437
            **kwargs,
Woosuk Kwon's avatar
Woosuk Kwon committed
438
439
440
441
442
443
        )

    def generate(
        self,
        prompts: List[str],
        sampling_params: SamplingParams,
444
        images: Optional[List[Image.Image]] = None,
445
    ) -> List[Tuple[List[List[int]], List[str]]]:
446
        if images is not None:
447
            assert len(prompts) == len(images)
448

449
450
451
        inputs = [TextPrompt(prompt=prompt) for prompt in prompts]
        if images is not None:
            for i, image in enumerate(images):
452
                inputs[i]["multi_modal_data"] = {"image": image}
453

454
        req_outputs = self.model.generate(inputs,
455
                                          sampling_params=sampling_params)
456
457

        outputs: List[Tuple[List[List[int]], List[str]]] = []
Woosuk Kwon's avatar
Woosuk Kwon committed
458
459
460
        for req_output in req_outputs:
            prompt_str = req_output.prompt
            prompt_ids = req_output.prompt_token_ids
461
462
            req_sample_output_ids: List[List[int]] = []
            req_sample_output_strs: List[str] = []
463
464
            for sample in req_output.outputs:
                output_str = sample.text
465
                output_ids = list(sample.token_ids)
466
467
468
                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
469
470
        return outputs

471
472
473
474
    def generate_w_logprobs(
        self,
        prompts: List[str],
        sampling_params: SamplingParams,
475
        images: Optional[List[Image.Image]] = None,
476
    ) -> List[Tuple[List[int], str, Optional[SampleLogprobs]]]:
477
478
        assert sampling_params.logprobs is not None

479
480
481
482
483
484
485
486
487
        if images is not None:
            assert len(prompts) == len(images)

        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": image}

        req_outputs = self.model.generate(inputs,
488
                                          sampling_params=sampling_params)
489
        outputs: List[Tuple[List[int], str, Optional[SampleLogprobs]]] = []
490
491
492
493
494
495
496
497
        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
498
499
500
501
    def generate_greedy(
        self,
        prompts: List[str],
        max_tokens: int,
502
        images: Optional[List[Image.Image]] = None,
Woosuk Kwon's avatar
Woosuk Kwon committed
503
504
    ) -> List[Tuple[List[int], str]]:
        greedy_params = SamplingParams(temperature=0.0, max_tokens=max_tokens)
505
        outputs = self.generate(prompts, greedy_params, images=images)
506
507
        return [(output_ids[0], output_str[0])
                for output_ids, output_str in outputs]
508

509
510
511
512
513
    def generate_greedy_logprobs(
        self,
        prompts: List[str],
        max_tokens: int,
        num_logprobs: int,
514
        images: Optional[List[Image.Image]] = None,
515
    ) -> List[Tuple[List[int], str, Optional[SampleLogprobs]]]:
516
517
518
        greedy_logprobs_params = SamplingParams(temperature=0.0,
                                                max_tokens=max_tokens,
                                                logprobs=num_logprobs)
519
520
521
        outputs = self.generate_w_logprobs(prompts,
                                           greedy_logprobs_params,
                                           images=images)
522
523
524
525

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

526
527
528
529
530
    def generate_beam_search(
        self,
        prompts: List[str],
        beam_width: int,
        max_tokens: int,
531
    ) -> List[Tuple[List[List[int]], List[str]]]:
532
533
534
535
536
537
        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
538

539
540
541
542
543
544
545
546
    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

547
548
549
550
    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
551
552
553
        del self.model
        cleanup()

Woosuk Kwon's avatar
Woosuk Kwon committed
554

555
@pytest.fixture(scope="session")
Woosuk Kwon's avatar
Woosuk Kwon committed
556
557
def vllm_runner():
    return VllmRunner
558
559
560
561
562
563
564
565
566


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={})
567
568
569
570
    if isinstance(tokenizer_group_type, type):
        return TokenizerPoolConfig(pool_size=1,
                                   pool_type=tokenizer_group_type,
                                   extra_config={})
571
    raise ValueError(f"Unknown tokenizer_group_type: {tokenizer_group_type}")
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587


@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
588
589
590
591
592
593
594


@pytest.fixture(scope="session")
def num_gpus_available():
    """Get number of GPUs without initializing the CUDA context
    in current process."""

595
    return cuda_device_count_stateless()