From b50646e5effd7cb5884cd96fdff4c53c18521198 Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Sun, 17 May 2026 22:57:59 -0500 Subject: [PATCH] [ROCm][CI] Stabilize ROCm pooling and multimodal CI (#42909) Signed-off-by: Andreas Karatzas --- tests/models/language/pooling/test_gritlm.py | 12 ++++-- .../pooling/test_max_tokens_per_doc.py | 43 +++++++++++++------ .../multimodal/generation/test_qwen2_5_vl.py | 5 ++- .../models/transformers/base.py | 8 ++++ 4 files changed, 51 insertions(+), 17 deletions(-) diff --git a/tests/models/language/pooling/test_gritlm.py b/tests/models/language/pooling/test_gritlm.py index 5ff5073e869..b1296a64171 100644 --- a/tests/models/language/pooling/test_gritlm.py +++ b/tests/models/language/pooling/test_gritlm.py @@ -7,7 +7,7 @@ from scipy.spatial.distance import cosine from vllm import LLM, SamplingParams from vllm.config import ModelConfig -from ....utils import RemoteOpenAIServer +from ....utils import ROCM_ENV_OVERRIDES, ROCM_EXTRA_ARGS, RemoteOpenAIServer from .embed_utils import run_client_embeddings MODEL_NAME = "parasail-ai/GritLM-7B-vllm" @@ -126,9 +126,15 @@ def test_gritlm_offline_embedding(vllm_runner): async def test_gritlm_api_server_embedding(): queries, q_instruction, documents, d_instruction = get_test_data() - args = ["--runner", "pooling", "--max_model_len", str(MAX_MODEL_LEN)] + args = [ + "--runner", + "pooling", + "--max_model_len", + str(MAX_MODEL_LEN), + *ROCM_EXTRA_ARGS, + ] - with RemoteOpenAIServer(MODEL_NAME, args) as server: + with RemoteOpenAIServer(MODEL_NAME, args, env_dict=ROCM_ENV_OVERRIDES) as server: client_embedding = server.get_async_client() d_rep = await run_client_embeddings( diff --git a/tests/models/language/pooling/test_max_tokens_per_doc.py b/tests/models/language/pooling/test_max_tokens_per_doc.py index 7b4e1fd0349..2e5ab70d274 100644 --- a/tests/models/language/pooling/test_max_tokens_per_doc.py +++ b/tests/models/language/pooling/test_max_tokens_per_doc.py @@ -17,6 +17,7 @@ from vllm.entrypoints.pooling.scoring.protocol import RerankResponse os.environ["VLLM_LOGGING_LEVEL"] = "WARNING" TEMPLATE_DIR = str(VLLM_PATH / "examples/pooling/score/template") +ExpectedPromptTokens = int | tuple[int, ...] long_query = "What is the capital of France?" * 20 long_doc = "The capital of France is Paris. " * 20 @@ -26,10 +27,10 @@ long_doc = "The capital of France is Paris. " * 20 class TestConfig: model: str args: list[str] - without_truncated_prompt_tokens: int - with_max_tokens_per_query_prompt_tokens: int - with_max_tokens_per_doc_prompt_tokens: int - with_max_tokens_per_query_and_doc_prompt_tokens: int + without_truncated_prompt_tokens: ExpectedPromptTokens + with_max_tokens_per_query_prompt_tokens: ExpectedPromptTokens + with_max_tokens_per_doc_prompt_tokens: ExpectedPromptTokens + with_max_tokens_per_query_and_doc_prompt_tokens: ExpectedPromptTokens RERANK_CONFIGS = [ @@ -79,8 +80,10 @@ RERANK_CONFIGS = [ "512", "--trust-remote-code", ], - without_truncated_prompt_tokens=286, - with_max_tokens_per_query_prompt_tokens=156, + # This model has produced both prompt-token totals in CI/local cache; + # keep truncation checks exact while tolerating the boundary delta. + without_truncated_prompt_tokens=(285, 286), + with_max_tokens_per_query_prompt_tokens=(155, 156), with_max_tokens_per_doc_prompt_tokens=155, with_max_tokens_per_query_and_doc_prompt_tokens=25, ), @@ -115,6 +118,13 @@ RERANK_CONFIGS = [ ] +def assert_prompt_tokens(actual: int, expected: ExpectedPromptTokens) -> None: + if isinstance(expected, int): + assert actual == expected + else: + assert actual in expected + + @pytest.fixture(scope="module", params=RERANK_CONFIGS, ids=lambda c: c.model) def server(request): config: TestConfig = request.param @@ -136,7 +146,10 @@ def test_without_truncated(server): assert rerank.id is not None assert rerank.results is not None assert len(rerank.results) == 1 - assert rerank.usage.prompt_tokens == config.without_truncated_prompt_tokens + assert_prompt_tokens( + rerank.usage.prompt_tokens, + config.without_truncated_prompt_tokens, + ) def test_max_tokens_per_query(server): @@ -158,7 +171,10 @@ def test_max_tokens_per_query(server): assert rerank.id is not None assert rerank.results is not None assert len(rerank.results) == 1 - assert rerank.usage.prompt_tokens == config.with_max_tokens_per_query_prompt_tokens + assert_prompt_tokens( + rerank.usage.prompt_tokens, + config.with_max_tokens_per_query_prompt_tokens, + ) def test_max_tokens_per_doc(server): @@ -180,7 +196,10 @@ def test_max_tokens_per_doc(server): assert rerank.id is not None assert rerank.results is not None assert len(rerank.results) == 1 - assert rerank.usage.prompt_tokens == config.with_max_tokens_per_doc_prompt_tokens + assert_prompt_tokens( + rerank.usage.prompt_tokens, + config.with_max_tokens_per_doc_prompt_tokens, + ) def test_max_tokens_per_query_and_doc(server): @@ -203,7 +222,7 @@ def test_max_tokens_per_query_and_doc(server): assert rerank.id is not None assert rerank.results is not None assert len(rerank.results) == 1 - assert ( - rerank.usage.prompt_tokens - == config.with_max_tokens_per_query_and_doc_prompt_tokens + assert_prompt_tokens( + rerank.usage.prompt_tokens, + config.with_max_tokens_per_query_and_doc_prompt_tokens, ) diff --git a/tests/models/multimodal/generation/test_qwen2_5_vl.py b/tests/models/multimodal/generation/test_qwen2_5_vl.py index 791bb3b3088..15a14da24d1 100644 --- a/tests/models/multimodal/generation/test_qwen2_5_vl.py +++ b/tests/models/multimodal/generation/test_qwen2_5_vl.py @@ -34,6 +34,7 @@ WINDOW_ATTN_IMAGE_PROMPT = qwen2_5_vl_chat_template( IMAGE_PLACEHOLDER, "Describe the image.", ) +IMAGE_ONLY_LIMIT_MM_PER_PROMPT = {"image": 1, "video": 0} def _window_attention_regression_image(): @@ -193,7 +194,7 @@ def test_qwen2_5_vl_window_attention_image( runner="generate", max_model_len=4096, dtype=dtype, - limit_mm_per_prompt={"image": 1}, + limit_mm_per_prompt=IMAGE_ONLY_LIMIT_MM_PER_PROMPT, compilation_config=_encoder_cudagraph_config(max_vision_items=1), ) as vllm_model: outputs = vllm_model.generate_greedy(prompt, max_tokens, images=images) @@ -231,7 +232,7 @@ def test_qwen2_5_vl_window_attention_image_batch( max_model_len=4096, max_num_seqs=2, dtype=dtype, - limit_mm_per_prompt={"image": 1}, + limit_mm_per_prompt=IMAGE_ONLY_LIMIT_MM_PER_PROMPT, compilation_config=_encoder_cudagraph_config(max_vision_items=2), ) as vllm_model: outputs = vllm_model.generate_greedy(prompts, max_tokens, images=images) diff --git a/vllm/model_executor/models/transformers/base.py b/vllm/model_executor/models/transformers/base.py index a3e4b844b80..e3591f2939c 100644 --- a/vllm/model_executor/models/transformers/base.py +++ b/vllm/model_executor/models/transformers/base.py @@ -353,6 +353,14 @@ class Base( for source, target in ccm.items(): orig_to_new_regex[re.compile(source)] = target + # Gemma3 checkpoints saved with older Transformers versions include an + # extra `vision_model` level that the current AutoModel no longer has. + vision_tower = getattr(self.model, "vision_tower", None) + if vision_tower is not None and not hasattr(vision_tower, "vision_model"): + orig_to_new_regex[ + re.compile(r"^(?:model\.)?vision_tower\.vision_model\.(.+)") + ] = r"model.vision_tower.\1" + # Handle unexpected weights which should be ignored if self.model._keys_to_ignore_on_load_unexpected is not None: for key in self.model._keys_to_ignore_on_load_unexpected: