| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import pytest |
| 5 | import torch.nn.functional as F |
| 6 | from transformers import AutoModelForImageTextToText |
| 7 | |
| 8 | from vllm.platforms import current_platform |
| 9 | |
| 10 | from ....conftest import IMAGE_ASSETS, HfRunner, PromptImageInput, VllmRunner |
| 11 | from ....utils import large_gpu_test |
| 12 | from ...utils import check_embeddings_close |
| 13 | |
| 14 | # Llava Next embedding implementation is only supported by CUDA. |
| 15 | # If run on ROCm, hf_model.model.resize_token_embeddings will |
| 16 | # cause the following error: |
| 17 | # RuntimeError: Calling torch.linalg.cholesky on a CUDA tensor |
| 18 | # requires compiling PyTorch with MAGMA. Please use PyTorch |
| 19 | # built with MAGMA support. |
| 20 | # If run on CPU, hf_model.model.resize_token_embeddings will |
| 21 | # cause the following error: |
| 22 | # RuntimeError: Calling torch.linalg.cholesky on a CPU tensor |
| 23 | # requires compiling PyTorch with LAPACK. Please use PyTorch |
| 24 | # built with LAPACK support. |
| 25 | pytestmark = pytest.mark.skipif( |
| 26 | not current_platform.is_cuda(), |
| 27 | reason="Llava Next model uses op that is only supported in CUDA", |
| 28 | ) |
| 29 | |
| 30 | llama3_template = "<|start_header_id|>user<|end_header_id|>\n\n{}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n \n" # noqa: E501 |
| 31 | |
| 32 | HF_TEXT_PROMPTS = [ |
| 33 | # T -> X |
| 34 | llama3_template.format( |
| 35 | "The label of the object is stop sign\nSummary above sentence in one word: " # noqa: E501 |
| 36 | ), |
| 37 | # T -> X |
| 38 | llama3_template.format("cherry blossom\nSummary above sentence in one word: "), |
| 39 | ] |
| 40 | |
| 41 | HF_IMAGE_PROMPTS = IMAGE_ASSETS.prompts( |
| 42 | { |
| 43 | # I -> X |
| 44 | "stop_sign": llama3_template.format( |
| 45 | "<image>\nSummary above image in one word: " |
| 46 | ), |
| 47 | # I -> X |
| 48 | "cherry_blossom": llama3_template.format( |
| 49 | "<image>\nSummary above image in one word: " |
| 50 | ), |
| 51 | } |
| 52 | ) |
| 53 | |
| 54 | MODELS = ["royokong/e5-v"] |
| 55 | |
| 56 | |
| 57 | def _run_test( |
| 58 | hf_runner: type[HfRunner], |
| 59 | vllm_runner: type[VllmRunner], |
| 60 | input_texts: list[str], |
| 61 | input_images: PromptImageInput, |
| 62 | model: str, |
| 63 | *, |
| 64 | dtype: str, |
| 65 | ) -> None: |
| 66 | # NOTE: take care of the order. run vLLM first, and then run HF. |
| 67 | # vLLM needs a fresh new process without cuda initialization. |
| 68 | # if we run HF first, the cuda initialization will be done and it |
| 69 | # will hurt multiprocessing backend with fork method (the default method). |
| 70 | with vllm_runner( |
| 71 | model, runner="pooling", dtype=dtype, max_model_len=4096, enforce_eager=True |
| 72 | ) as vllm_model: |
| 73 | vllm_outputs = vllm_model.embed(input_texts, images=input_images) |
| 74 | |
| 75 | with hf_runner( |
| 76 | model, dtype=dtype, auto_cls=AutoModelForImageTextToText |
| 77 | ) as hf_model: |
| 78 | # Patch the issue where generation_config.json is missing |
| 79 | hf_model.processor.patch_size = hf_model.model.config.vision_config.patch_size |
| 80 | |
| 81 | # Patch the issue where image_token_id |
| 82 | # exceeds the maximum allowed vocab size |
| 83 | hf_model.model.resize_token_embeddings( |
| 84 | hf_model.model.model.language_model.vocab_size + 1 |
| 85 | ) |
| 86 | |
| 87 | all_inputs = hf_model.get_inputs(input_texts, images=input_images) |
| 88 | |
| 89 | all_outputs = [] |
| 90 | for inputs in all_inputs: |
| 91 | # Based on: https://huggingface.co/royokong/e5-v |
| 92 | outputs = hf_model.model( |
| 93 | **hf_model.wrap_device(inputs), |
| 94 | return_dict=True, |
| 95 | output_hidden_states=True, |
| 96 | ) |
| 97 | pooled_output = F.normalize(outputs.hidden_states[-1][0, -1, :], dim=-1) |
| 98 | |
| 99 | all_outputs.append(pooled_output.tolist()) |
| 100 | |
| 101 | hf_outputs = all_outputs |
| 102 | |
| 103 | check_embeddings_close( |
| 104 | embeddings_0_lst=hf_outputs, |
| 105 | embeddings_1_lst=vllm_outputs, |
| 106 | name_0="hf", |
| 107 | name_1="vllm", |
| 108 | ) |
| 109 | |
| 110 | |
| 111 | @pytest.mark.core_model |
| 112 | @pytest.mark.parametrize("model", MODELS) |
| 113 | @pytest.mark.parametrize("dtype", ["half"]) |
| 114 | def test_models_text( |
| 115 | hf_runner, |
| 116 | vllm_runner, |
| 117 | image_assets, |
| 118 | model: str, |
| 119 | dtype: str, |
| 120 | ) -> None: |
| 121 | input_texts_images = [(text, None) for text in HF_TEXT_PROMPTS] |
| 122 | input_texts = [text for text, _ in input_texts_images] |
| 123 | input_images = [image for _, image in input_texts_images] |
| 124 | |
| 125 | _run_test( |
| 126 | hf_runner, |
| 127 | vllm_runner, |
| 128 | input_texts, |
| 129 | input_images, # type: ignore |
| 130 | model, |
| 131 | dtype=dtype, |
| 132 | ) |
| 133 | |
| 134 | |
| 135 | @large_gpu_test(min_gb=48) |
| 136 | @pytest.mark.core_model |
| 137 | @pytest.mark.parametrize("model", MODELS) |
| 138 | @pytest.mark.parametrize("dtype", ["half"]) |
| 139 | def test_models_image( |
| 140 | hf_runner, |
| 141 | vllm_runner, |
| 142 | image_assets, |
| 143 | model: str, |
| 144 | dtype: str, |
| 145 | ) -> None: |
| 146 | input_texts_images = [ |
| 147 | (text, asset.pil_image) for text, asset in zip(HF_IMAGE_PROMPTS, image_assets) |
| 148 | ] |
| 149 | input_texts = [text for text, _ in input_texts_images] |
| 150 | input_images = [image for _, image in input_texts_images] |
| 151 | |
| 152 | _run_test( |
| 153 | hf_runner, |
| 154 | vllm_runner, |
| 155 | input_texts, |
| 156 | input_images, |
| 157 | model, |
| 158 | dtype=dtype, |
| 159 | ) |