@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 129 lines 3.94 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 # ruff: noqa: SIM117
4
5 from functools import partial
6
7 import pytest
8
9 from ...utils import EmbedModelInfo
10
11 MODELS = [
12 EmbedModelInfo(
13 "nomic-ai/nomic-embed-text-v1",
14 # Fixme:
15 # Update nomic-embed code to support the latest
16 # HF version and remove revision set.
17 revision="720244025c1a7e15661a174c63cce63c8218e52b",
18 ),
19 # EmbedModelInfo("nomic-ai/nomic-embed-text-v1.5"),
20 # EmbedModelInfo("nomic-ai/CodeRankEmbed"),
21 EmbedModelInfo("nomic-ai/nomic-embed-text-v2-moe"),
22 # EmbedModelInfo("Snowflake/snowflake-arctic-embed-m-long"),
23 ]
24
25 rope_theta = 1000
26 factor = 4.0
27 original_max_position_embeddings = 2048
28 max_model_len = int(original_max_position_embeddings * factor)
29
30
31 @pytest.fixture(scope="module", autouse=True)
32 def predownload_models():
33 """Warm the HF cache so engine startup never hits transient hub errors."""
34 from vllm.transformers_utils.repo_utils import hf_api, with_retry
35
36 for model_info in MODELS:
37 with_retry(
38 partial(
39 hf_api().snapshot_download,
40 model_info.name,
41 revision=model_info.revision,
42 ),
43 f"Error downloading {model_info.name}",
44 )
45
46
47 @pytest.mark.parametrize("model_info", MODELS)
48 def test_default(model_info, vllm_runner):
49 with vllm_runner(
50 model_info.name,
51 revision=model_info.revision,
52 runner="pooling",
53 max_model_len=None,
54 ) as vllm_model:
55 model_config = vllm_model.llm.llm_engine.model_config
56 if model_info.name == "nomic-ai/nomic-embed-text-v2-moe":
57 # For nomic-embed-text-v2-moe the length is set to 512
58 # by sentence_bert_config.json.
59 assert model_config.max_model_len == 512
60 if model_info.name == "nomic-ai/nomic-embed-text-v1":
61 assert model_config.max_model_len == 8192
62
63
64 @pytest.mark.parametrize("model_info", MODELS)
65 def test_set_max_model_len_legal(model_info, vllm_runner):
66 # set max_model_len <= 512
67 with vllm_runner(
68 model_info.name,
69 revision=model_info.revision,
70 runner="pooling",
71 max_model_len=256,
72 ) as vllm_model:
73 model_config = vllm_model.llm.llm_engine.model_config
74 assert model_config.max_model_len == 256
75
76 # For nomic-embed-text-v2-moe the length is set to 512
77 # by sentence_bert_config.json.
78 if model_info.name == "nomic-ai/nomic-embed-text-v2-moe":
79 with pytest.raises(ValueError):
80 with vllm_runner(
81 model_info.name,
82 revision=model_info.revision,
83 runner="pooling",
84 max_model_len=1024,
85 ):
86 pass
87 return
88
89 # set 512 < max_model_len <= 2048
90 with vllm_runner(
91 model_info.name,
92 revision=model_info.revision,
93 runner="pooling",
94 max_model_len=1024,
95 ) as vllm_model:
96 model_config = vllm_model.llm.llm_engine.model_config
97 assert model_config.max_model_len == 1024
98
99 # set max_model_len > 2048
100 with vllm_runner(
101 model_info.name,
102 revision=model_info.revision,
103 runner="pooling",
104 max_model_len=4096,
105 ) as vllm_model:
106 model_config = vllm_model.llm.llm_engine.model_config
107 assert model_config.max_model_len == 4096
108
109
110 @pytest.mark.parametrize("model_info", MODELS)
111 def test_use_rope_scaling_legal(model_info, vllm_runner):
112 hf_overrides = {
113 "rope_parameters": {
114 "rope_theta": rope_theta,
115 "rope_type": "yarn",
116 "factor": factor,
117 "original_max_position_embeddings": original_max_position_embeddings,
118 },
119 "max_model_len": max_model_len,
120 }
121
122 with vllm_runner(
123 model_info.name,
124 revision=model_info.revision,
125 runner="pooling",
126 max_model_len=None,
127 hf_overrides=hf_overrides,
128 ):
129 pass