conftest.py 18.2 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
17
from vllm.distributed import (destroy_distributed_environment,
                              destroy_model_parallel)
18
from vllm.inputs import TextPrompt
19
from vllm.logger import init_logger
20
21
22
from vllm.multimodal import MultiModalData
from vllm.multimodal.image import ImageFeatureData, ImagePixelData
from vllm.sequence import SampleLogprobs
23
from vllm.utils import cuda_device_count_stateless, is_cpu
24
25

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

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

31
# Multi modal related
32
# You can use `.buildkite/download-images.sh` to download the assets
33
PIXEL_VALUES_FILES = [
34
35
36
    os.path.join(_TEST_DIR, "images", filename) for filename in
    ["stop_sign_pixel_values.pt", "cherry_blossom_pixel_values.pt"]
]
37
IMAGE_FEATURES_FILES = [
38
39
40
    os.path.join(_TEST_DIR, "images", filename) for filename in
    ["stop_sign_image_features.pt", "cherry_blossom_image_features.pt"]
]
41
IMAGE_FILES = [
42
43
44
    os.path.join(_TEST_DIR, "images", filename)
    for filename in ["stop_sign.jpg", "cherry_blossom.jpg"]
]
45
assert len(PIXEL_VALUES_FILES) == len(IMAGE_FEATURES_FILES) == len(IMAGE_FILES)
46

47

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


54
55
def cleanup():
    destroy_model_parallel()
56
    destroy_distributed_environment()
57
58
59
    with contextlib.suppress(AssertionError):
        torch.distributed.destroy_process_group()
    gc.collect()
60
61
    if not is_cpu():
        torch.cuda.empty_cache()
62
63


64
@pytest.fixture()
65
def should_do_global_cleanup_after_test(request) -> bool:
66
67
68
69
    """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.
    """
70
71
72
73

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

74
75
76
    return True


77
@pytest.fixture(autouse=True)
78
def cleanup_fixture(should_do_global_cleanup_after_test: bool):
79
    yield
80
81
    if should_do_global_cleanup_after_test:
        cleanup()
82
83


84
85
@pytest.fixture(scope="session")
def hf_images() -> List[Image.Image]:
86
    return [Image.open(filename) for filename in IMAGE_FILES]
87
88
89


@pytest.fixture()
90
def vllm_images(request) -> List[MultiModalData]:
91
92
93
    vision_language_config = request.getfixturevalue("model_and_config")[1]
    if vision_language_config.image_input_type == (
            VisionLanguageConfig.ImageInputType.IMAGE_FEATURES):
94
95
        return [
            ImageFeatureData(torch.load(filename))
96
            for filename in IMAGE_FEATURES_FILES
97
        ]
98
    else:
99
        return [
100
            ImagePixelData(Image.open(filename)) for filename in IMAGE_FILES
101
102
103
104
105
        ]


@pytest.fixture()
def vllm_image_tensors(request) -> List[torch.Tensor]:
106
    return [torch.load(filename) for filename in PIXEL_VALUES_FILES]
107
108


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


@pytest.fixture
def example_long_prompts() -> List[str]:
    prompts = []
    for filename in _LONG_PROMPTS:
121
        prompts += _read_prompts(filename)
122
    return prompts
Woosuk Kwon's avatar
Woosuk Kwon committed
123
124
125
126
127
128
129
130


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

131
_T = TypeVar("_T", nn.Module, torch.Tensor, BatchEncoding)
132

Woosuk Kwon's avatar
Woosuk Kwon committed
133
134
135

class HfRunner:

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

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

154
        self.model_name = model_name
155

156
        if is_embedding_model:
157
158
            # Lazy init required for AMD CI
            from sentence_transformers import SentenceTransformer
159
160
161
162
163
            self.model = self.wrap_device(
                SentenceTransformer(
                    model_name,
                    device="cpu",
                ).to(dtype=torch_dtype))
164
        else:
165
166
167
168
169
            if is_vision_model:
                auto_cls = AutoModelForVision2Seq
            else:
                auto_cls = AutoModelForCausalLM

170
            model_kwargs = model_kwargs if model_kwargs is not None else {}
171
            self.model = self.wrap_device(
172
                auto_cls.from_pretrained(
173
174
175
                    model_name,
                    torch_dtype=torch_dtype,
                    trust_remote_code=True,
176
                    **model_kwargs,
177
                ))
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195

        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
196
197
198
199

    def generate(
        self,
        prompts: List[str],
200
        images: Optional[List[Image.Image]] = None,
Woosuk Kwon's avatar
Woosuk Kwon committed
201
        **kwargs,
202
    ) -> List[Tuple[List[List[int]], List[str]]]:
203
204
        if images:
            assert len(prompts) == len(images)
205
206

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

    def generate_greedy(
        self,
        prompts: List[str],
        max_tokens: int,
235
        images: Optional[List[Image.Image]] = None,
Chang Su's avatar
Chang Su committed
236
        **kwargs,
Woosuk Kwon's avatar
Woosuk Kwon committed
237
    ) -> List[Tuple[List[int], str]]:
238
239
        outputs = self.generate(prompts,
                                do_sample=False,
240
                                max_new_tokens=max_tokens,
Chang Su's avatar
Chang Su committed
241
242
                                images=images,
                                **kwargs)
243
244
245

        return [(output_ids[0], output_str[0])
                for output_ids, output_str in outputs]
246
247
248
249
250
251

    def generate_beam_search(
        self,
        prompts: List[str],
        beam_width: int,
        max_tokens: int,
252
    ) -> List[Tuple[List[List[int]], List[str]]]:
253
254
255
256
257
258
259
260
261
262
263
264
265
266
        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
267

268
269
270
271
272
273
274
275
276
    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(
277
                self.wrap_device(input_ids),
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
                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)
294
                logprobs = F.log_softmax(logits, dim=-1, dtype=torch.float32)
295
296
297
298
                seq_logprobs.append(logprobs)
            all_logprobs.append(seq_logprobs)
        return all_logprobs

299
300
301
302
303
    def generate_greedy_logprobs_limit(
        self,
        prompts: List[str],
        max_tokens: int,
        num_logprobs: int,
304
305
306
307
    ) -> 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] = []
308
309
310
311

        for prompt in prompts:
            input_ids = self.tokenizer(prompt, return_tensors="pt").input_ids
            output = self.model.generate(
312
                self.wrap_device(input_ids),
313
314
315
316
317
318
319
                use_cache=True,
                do_sample=False,
                max_new_tokens=max_tokens,
                output_hidden_states=True,
                return_dict_in_generate=True,
            )

320
            seq_logprobs: List[torch.Tensor] = []
321
322
323
324
325
326
327
328
329
330
            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)
331
                logprobs = F.log_softmax(logits, dim=-1, dtype=torch.float32)
332
333
334
                seq_logprobs.append(logprobs)

            # convert to dict
335
            seq_logprobs_lst: List[Dict[int, float]] = []
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
            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]

359
360
361
    def encode(self, prompts: List[str]) -> List[List[torch.Tensor]]:
        return self.model.encode(prompts)

362
363
364
365
    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
366
367
368
        del self.model
        cleanup()

Woosuk Kwon's avatar
Woosuk Kwon committed
369

Cyrus Leung's avatar
Cyrus Leung committed
370
@pytest.fixture(scope="session")
Woosuk Kwon's avatar
Woosuk Kwon committed
371
372
373
374
375
376
377
378
379
380
def hf_runner():
    return HfRunner


class VllmRunner:

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

    def generate(
        self,
        prompts: List[str],
        sampling_params: SamplingParams,
412
        images: Optional[List[MultiModalData]] = None,
413
    ) -> List[Tuple[List[List[int]], List[str]]]:
414
        if images is not None:
415
            assert len(prompts) == len(images)
416

417
418
419
420
        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
421

422
        req_outputs = self.model.generate(inputs,
423
                                          sampling_params=sampling_params)
424
425

        outputs: List[Tuple[List[List[int]], List[str]]] = []
Woosuk Kwon's avatar
Woosuk Kwon committed
426
427
428
        for req_output in req_outputs:
            prompt_str = req_output.prompt
            prompt_ids = req_output.prompt_token_ids
429
430
            req_sample_output_ids: List[List[int]] = []
            req_sample_output_strs: List[str] = []
431
432
433
434
435
436
            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
437
438
        return outputs

439
440
441
442
    def generate_w_logprobs(
        self,
        prompts: List[str],
        sampling_params: SamplingParams,
443
    ) -> List[Tuple[List[int], str, Optional[SampleLogprobs]]]:
444
445
446
447
        assert sampling_params.logprobs is not None

        req_outputs = self.model.generate(prompts,
                                          sampling_params=sampling_params)
448
        outputs: List[Tuple[List[int], str, Optional[SampleLogprobs]]] = []
449
450
451
452
453
454
455
456
        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
457
458
459
460
    def generate_greedy(
        self,
        prompts: List[str],
        max_tokens: int,
461
        images: Optional[List[MultiModalData]] = None,
Woosuk Kwon's avatar
Woosuk Kwon committed
462
463
    ) -> List[Tuple[List[int], str]]:
        greedy_params = SamplingParams(temperature=0.0, max_tokens=max_tokens)
464
        outputs = self.generate(prompts, greedy_params, images=images)
465
466
        return [(output_ids[0], output_str[0])
                for output_ids, output_str in outputs]
467

468
469
470
471
472
    def generate_greedy_logprobs(
        self,
        prompts: List[str],
        max_tokens: int,
        num_logprobs: int,
473
    ) -> List[Tuple[List[int], str, Optional[SampleLogprobs]]]:
474
475
476
477
478
479
480
481
        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]

482
483
484
485
486
    def generate_beam_search(
        self,
        prompts: List[str],
        beam_width: int,
        max_tokens: int,
487
    ) -> List[Tuple[List[List[int]], List[str]]]:
488
489
490
491
492
493
        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
494

495
496
497
498
499
500
501
502
    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

503
504
505
506
    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
507
508
509
        del self.model
        cleanup()

Woosuk Kwon's avatar
Woosuk Kwon committed
510

511
@pytest.fixture(scope="session")
Woosuk Kwon's avatar
Woosuk Kwon committed
512
513
def vllm_runner():
    return VllmRunner
514
515
516
517
518
519
520
521
522
523


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}")
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539


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


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

547
    return cuda_device_count_stateless()