| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import pytest |
| 5 | import torch |
| 6 | from transformers import AutoModel |
| 7 | |
| 8 | from vllm import PoolingParams |
| 9 | from vllm.config import PoolerConfig |
| 10 | |
| 11 | from ...utils import check_embeddings_close |
| 12 | |
| 13 | |
| 14 | @pytest.mark.parametrize( |
| 15 | "model", |
| 16 | [ |
| 17 | # Be careful of the order of models, decoder-only models should be |
| 18 | # placed before encoder-only models, otherwise `Qwen2.5-0.5B-Instruct` |
| 19 | # case won't pass because gte-Qwen2-1.5B-instruct will cache custom |
| 20 | # model code with bidirectional attention. |
| 21 | # [Decoder-only] |
| 22 | pytest.param( |
| 23 | "BAAI/bge-multilingual-gemma2", |
| 24 | marks=[pytest.mark.core_model, pytest.mark.slow_test], |
| 25 | ), |
| 26 | pytest.param( |
| 27 | "intfloat/e5-mistral-7b-instruct", |
| 28 | marks=[pytest.mark.core_model, pytest.mark.cpu_model], |
| 29 | ), |
| 30 | pytest.param( |
| 31 | "ssmits/Qwen2-7B-Instruct-embed-base", marks=[pytest.mark.cpu_model] |
| 32 | ), |
| 33 | # [Encoder-only] |
| 34 | pytest.param( |
| 35 | "BAAI/bge-base-en-v1.5", |
| 36 | marks=[ |
| 37 | pytest.mark.core_model, |
| 38 | pytest.mark.cpu_model, |
| 39 | pytest.mark.slow_test, |
| 40 | ], |
| 41 | ), |
| 42 | pytest.param("sentence-transformers/all-MiniLM-L12-v2"), |
| 43 | pytest.param("intfloat/multilingual-e5-small"), |
| 44 | # [Cross-Encoder] |
| 45 | pytest.param( |
| 46 | "sentence-transformers/stsb-roberta-base-v2", |
| 47 | marks=[pytest.mark.core_model, pytest.mark.cpu_model], |
| 48 | ), |
| 49 | ], |
| 50 | ) |
| 51 | def test_models( |
| 52 | hf_runner, |
| 53 | vllm_runner, |
| 54 | example_prompts, |
| 55 | model, |
| 56 | ) -> None: |
| 57 | vllm_extra_kwargs = {} |
| 58 | if model == "ssmits/Qwen2-7B-Instruct-embed-base": |
| 59 | vllm_extra_kwargs["pooler_config"] = PoolerConfig( |
| 60 | seq_pooling_type="MEAN", use_activation=False |
| 61 | ) |
| 62 | |
| 63 | max_model_len: int | None = 512 |
| 64 | if model in [ |
| 65 | "sentence-transformers/all-MiniLM-L12-v2", |
| 66 | "sentence-transformers/stsb-roberta-base-v2", |
| 67 | ]: |
| 68 | max_model_len = None |
| 69 | |
| 70 | # The example_prompts has ending "\n", for example: |
| 71 | # "Write a short story about a robot that dreams for the first time.\n" |
| 72 | # sentence_transformers will strip the input texts, see: |
| 73 | # https://github.com/UKPLab/sentence-transformers/blob/v3.1.1/sentence_transformers/models/Transformer.py#L159 |
| 74 | # This makes the input_ids different between hf_model and vllm_model. |
| 75 | # So we need to strip the input texts to avoid test failing. |
| 76 | example_prompts = [str(s).strip() for s in example_prompts] |
| 77 | |
| 78 | with hf_runner(model, is_sentence_transformer=True) as hf_model: |
| 79 | hf_outputs = hf_model.encode(example_prompts) |
| 80 | |
| 81 | with vllm_runner( |
| 82 | model, runner="pooling", max_model_len=max_model_len, **vllm_extra_kwargs |
| 83 | ) as vllm_model: |
| 84 | vllm_outputs = vllm_model.embed(example_prompts) |
| 85 | |
| 86 | check_embeddings_close( |
| 87 | embeddings_0_lst=hf_outputs, |
| 88 | embeddings_1_lst=vllm_outputs, |
| 89 | name_0="hf", |
| 90 | name_1="vllm", |
| 91 | tol=1e-2, |
| 92 | ) |
| 93 | |
| 94 | |
| 95 | @pytest.mark.parametrize( |
| 96 | "model", |
| 97 | [ |
| 98 | "BAAI/bge-base-en-v1.5", |
| 99 | "intfloat/multilingual-e5-small", |
| 100 | ], |
| 101 | ) |
| 102 | @torch.inference_mode() |
| 103 | def test_encoder_only_model_runner_v2_attention( |
| 104 | hf_runner, |
| 105 | vllm_runner, |
| 106 | monkeypatch, |
| 107 | model: str, |
| 108 | ) -> None: |
| 109 | prompts = [ |
| 110 | "short input", |
| 111 | "a longer input that exercises mixed sequence lengths", |
| 112 | ] |
| 113 | |
| 114 | with hf_runner(model, dtype="float", auto_cls=AutoModel) as hf_model: |
| 115 | hf_outputs = [] |
| 116 | for prompt in prompts: |
| 117 | inputs = hf_model.tokenizer(prompt, return_tensors="pt") |
| 118 | output = hf_model.model(**hf_model.wrap_device(inputs)) |
| 119 | embedding = torch.nn.functional.normalize( |
| 120 | output.last_hidden_state[0, -1].float(), dim=0 |
| 121 | ) |
| 122 | hf_outputs.append(embedding.cpu().tolist()) |
| 123 | |
| 124 | monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1") |
| 125 | with vllm_runner( |
| 126 | model, |
| 127 | runner="pooling", |
| 128 | dtype="float", |
| 129 | max_model_len=64, |
| 130 | max_num_seqs=2, |
| 131 | gpu_memory_utilization=0.25, |
| 132 | pooler_config=PoolerConfig( |
| 133 | task="embed", seq_pooling_type="LAST", use_activation=True |
| 134 | ), |
| 135 | ) as vllm_model: |
| 136 | assert vllm_model.llm.llm_engine.vllm_config.use_v2_model_runner |
| 137 | vllm_outputs = vllm_model.embed(prompts) |
| 138 | |
| 139 | check_embeddings_close( |
| 140 | embeddings_0_lst=hf_outputs, |
| 141 | embeddings_1_lst=vllm_outputs, |
| 142 | name_0="hf", |
| 143 | name_1="vllm", |
| 144 | tol=1e-2, |
| 145 | ) |
| 146 | |
| 147 | |
| 148 | @pytest.mark.core_model |
| 149 | def test_encoder_model_runner_v2(hf_runner, vllm_runner, monkeypatch) -> None: |
| 150 | model = "sentence-transformers/all-MiniLM-L6-v2" |
| 151 | prompt_batches = [ |
| 152 | ["short input"], |
| 153 | [ |
| 154 | "short input", |
| 155 | "a longer input that exercises mixed sequence lengths", |
| 156 | ], |
| 157 | ] |
| 158 | |
| 159 | with hf_runner(model, is_sentence_transformer=True) as hf_model: |
| 160 | hf_outputs = [hf_model.encode(prompts) for prompts in prompt_batches] |
| 161 | |
| 162 | monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1") |
| 163 | with vllm_runner( |
| 164 | model, |
| 165 | runner="pooling", |
| 166 | max_model_len=64, |
| 167 | ) as vllm_model: |
| 168 | assert vllm_model.llm.llm_engine.vllm_config.use_v2_model_runner |
| 169 | vllm_outputs = [vllm_model.embed(prompts) for prompts in prompt_batches] |
| 170 | |
| 171 | for hf_batch, vllm_batch in zip(hf_outputs, vllm_outputs): |
| 172 | check_embeddings_close( |
| 173 | embeddings_0_lst=hf_batch, |
| 174 | embeddings_1_lst=vllm_batch, |
| 175 | name_0="hf", |
| 176 | name_1="vllm", |
| 177 | tol=1e-2, |
| 178 | ) |
| 179 | |
| 180 | |
| 181 | @pytest.mark.core_model |
| 182 | def test_matryoshka_dimensions_model_runner_v2( |
| 183 | hf_runner, vllm_runner, monkeypatch |
| 184 | ) -> None: |
| 185 | model = "Snowflake/snowflake-arctic-embed-m-v1.5" |
| 186 | prompts = ["short input", "a longer input for a different output width"] |
| 187 | dimensions = [None, 256] |
| 188 | |
| 189 | with hf_runner(model, is_sentence_transformer=True) as hf_model: |
| 190 | hf_outputs = hf_model.encode(prompts) |
| 191 | |
| 192 | monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1") |
| 193 | with vllm_runner( |
| 194 | model, |
| 195 | runner="pooling", |
| 196 | max_model_len=64, |
| 197 | gpu_memory_utilization=0.25, |
| 198 | ) as vllm_model: |
| 199 | assert vllm_model.llm.llm_engine.vllm_config.use_v2_model_runner |
| 200 | vllm_outputs = vllm_model.embed( |
| 201 | prompts, |
| 202 | pooling_params=[PoolingParams(dimensions=d) for d in dimensions], |
| 203 | ) |
| 204 | |
| 205 | expected_outputs = [] |
| 206 | for output, dimension in zip(hf_outputs, dimensions): |
| 207 | output = torch.as_tensor(output) |
| 208 | if dimension is not None: |
| 209 | output = torch.nn.functional.normalize(output[:dimension], dim=0) |
| 210 | expected_outputs.append(output.tolist()) |
| 211 | |
| 212 | assert [len(output) for output in vllm_outputs] == [768, 256] |
| 213 | check_embeddings_close( |
| 214 | embeddings_0_lst=expected_outputs, |
| 215 | embeddings_1_lst=vllm_outputs, |
| 216 | name_0="hf", |
| 217 | name_1="vllm", |
| 218 | tol=1e-2, |
| 219 | ) |