| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import os |
| 5 | |
| 6 | import pytest |
| 7 | |
| 8 | from vllm.model_executor.layers.pooler import DispatchPooler |
| 9 | from vllm.model_executor.layers.pooler.seqwise import CLSPool, MeanPool |
| 10 | from vllm.model_executor.models.bert import BertEmbeddingModel |
| 11 | from vllm.model_executor.models.roberta import RobertaEmbeddingModel |
| 12 | from vllm.platforms import current_platform |
| 13 | |
| 14 | MAX_MODEL_LEN = 128 |
| 15 | MODEL_NAME = os.environ.get("MODEL_NAME", "BAAI/bge-base-en-v1.5") |
| 16 | REVISION = os.environ.get("REVISION", "main") |
| 17 | |
| 18 | MODEL_NAME_ROBERTA = os.environ.get("MODEL_NAME", "intfloat/multilingual-e5-base") |
| 19 | REVISION_ROBERTA = os.environ.get("REVISION", "main") |
| 20 | |
| 21 | |
| 22 | @pytest.mark.skipif( |
| 23 | current_platform.is_rocm(), reason="Xformers backend is not supported on ROCm." |
| 24 | ) |
| 25 | def test_model_loading_with_params(vllm_runner, monkeypatch): |
| 26 | """Test parameter weight loading with tp>1.""" |
| 27 | # to use apply_model |
| 28 | monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1") |
| 29 | with vllm_runner( |
| 30 | model_name=MODEL_NAME, |
| 31 | revision=REVISION, |
| 32 | dtype="float16", |
| 33 | max_model_len=MAX_MODEL_LEN, |
| 34 | ) as vllm_model: |
| 35 | output = vllm_model.embed( |
| 36 | "Write a short story about a robot that dreams for the first time.\n" |
| 37 | ) |
| 38 | |
| 39 | model_config = vllm_model.llm.llm_engine.model_config |
| 40 | model_tokenizer = vllm_model.llm.llm_engine.tokenizer |
| 41 | |
| 42 | # asserts on the bert model config file |
| 43 | assert model_config.encoder_config["max_seq_length"] == 512 |
| 44 | assert model_config.encoder_config["do_lower_case"] |
| 45 | |
| 46 | # asserts on the pooling config files |
| 47 | assert model_config.pooler_config.seq_pooling_type == "CLS" |
| 48 | assert model_config.pooler_config.tok_pooling_type == "ALL" |
| 49 | assert model_config.pooler_config.use_activation |
| 50 | |
| 51 | # asserts on the tokenizer loaded |
| 52 | assert model_config.tokenizer == "BAAI/bge-base-en-v1.5" |
| 53 | assert model_tokenizer.model_max_length == 512 |
| 54 | |
| 55 | def check_model(model): |
| 56 | assert isinstance(model, BertEmbeddingModel) |
| 57 | assert isinstance(pooler := model.pooler, DispatchPooler) |
| 58 | assert isinstance(pooler.poolers_by_task["embed"].pooling, CLSPool) |
| 59 | |
| 60 | vllm_model.apply_model(check_model) |
| 61 | |
| 62 | assert output |
| 63 | |
| 64 | |
| 65 | @pytest.mark.skipif( |
| 66 | current_platform.is_rocm(), reason="Xformers backend is not supported on ROCm." |
| 67 | ) |
| 68 | def test_roberta_model_loading_with_params(vllm_runner, monkeypatch): |
| 69 | """Test parameter weight loading with tp>1.""" |
| 70 | # to use apply_model |
| 71 | monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1") |
| 72 | with vllm_runner( |
| 73 | model_name=MODEL_NAME_ROBERTA, |
| 74 | revision=REVISION_ROBERTA, |
| 75 | dtype="float16", |
| 76 | max_model_len=MAX_MODEL_LEN, |
| 77 | ) as vllm_model: |
| 78 | output = vllm_model.embed( |
| 79 | "Write a short story about a robot that dreams for the first time.\n" |
| 80 | ) |
| 81 | |
| 82 | model_config = vllm_model.llm.llm_engine.model_config |
| 83 | model_tokenizer = vllm_model.llm.llm_engine.tokenizer |
| 84 | |
| 85 | # asserts on the bert model config file |
| 86 | assert model_config.encoder_config["max_seq_length"] == 512 |
| 87 | assert not model_config.encoder_config["do_lower_case"] |
| 88 | |
| 89 | # asserts on the pooling config files |
| 90 | assert model_config.pooler_config.seq_pooling_type == "MEAN" |
| 91 | assert model_config.pooler_config.tok_pooling_type == "ALL" |
| 92 | assert model_config.pooler_config.use_activation |
| 93 | |
| 94 | # asserts on the tokenizer loaded |
| 95 | assert model_config.tokenizer == "intfloat/multilingual-e5-base" |
| 96 | assert model_tokenizer.model_max_length == 512 |
| 97 | |
| 98 | def check_model(model): |
| 99 | assert isinstance(model, RobertaEmbeddingModel) |
| 100 | assert isinstance(pooler := model.pooler, DispatchPooler) |
| 101 | assert isinstance(pooler.poolers_by_task["embed"].pooling, MeanPool) |
| 102 | |
| 103 | vllm_model.apply_model(check_model) |
| 104 | |
| 105 | assert output |
| 106 | |
| 107 | |
| 108 | @pytest.mark.skipif( |
| 109 | current_platform.is_rocm(), reason="Xformers backend is not supported on ROCm." |
| 110 | ) |
| 111 | def test_facebook_roberta_model_loading_with_params(vllm_runner, monkeypatch): |
| 112 | """Test loading roberta-base model with no lm_head.""" |
| 113 | # to use apply_model |
| 114 | monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1") |
| 115 | model_name = "FacebookAI/roberta-base" |
| 116 | with vllm_runner( |
| 117 | model_name=model_name, dtype="float16", max_model_len=MAX_MODEL_LEN |
| 118 | ) as vllm_model: |
| 119 | output = vllm_model.embed( |
| 120 | "Write a short story about a robot that dreams for the first time.\n" |
| 121 | ) |
| 122 | |
| 123 | assert vllm_model.llm.llm_engine.model_config.tokenizer == model_name |
| 124 | |
| 125 | def check_model(model): |
| 126 | assert isinstance(model, RobertaEmbeddingModel) |
| 127 | assert not hasattr(model, "lm_head") |
| 128 | assert isinstance(pooler := model.pooler, DispatchPooler) |
| 129 | assert isinstance(pooler.poolers_by_task["embed"].pooling, CLSPool) |
| 130 | |
| 131 | vllm_model.apply_model(check_model) |
| 132 | |
| 133 | assert output |