From 47672f38b58581cf2b7c33201e6ae01639c5ff51 Mon Sep 17 00:00:00 2001
From: Isotr0py <2037008807@qq.com>
Date: Sat, 9 Nov 2024 12:02:59 +0800
Subject: [PATCH] [CI/Build] Fix VLM broadcast tests `tensor_parallel_size`
 passing (#10161)

Signed-off-by: Isotr0py <2037008807@qq.com>
---
 tests/models/decoder_only/vision_language/test_models.py     | 1 +
 tests/models/decoder_only/vision_language/vlm_utils/types.py | 1 +
 2 files changed, 2 insertions(+)

diff --git a/tests/models/decoder_only/vision_language/test_models.py b/tests/models/decoder_only/vision_language/test_models.py
index 163752e9..1ab42f8c 100644
--- a/tests/models/decoder_only/vision_language/test_models.py
+++ b/tests/models/decoder_only/vision_language/test_models.py
@@ -34,6 +34,7 @@ COMMON_BROADCAST_SETTINGS = {
     "dtype": "half",
     "max_tokens": 5,
     "tensor_parallel_size": 2,
+    "model_kwargs": {"device_map": "auto"},
     "image_size_factors": [(.25, 0.5, 1.0)],
     "distributed_executor_backend": (
         "ray",
diff --git a/tests/models/decoder_only/vision_language/vlm_utils/types.py b/tests/models/decoder_only/vision_language/vlm_utils/types.py
index fd18c7c8..8459476d 100644
--- a/tests/models/decoder_only/vision_language/vlm_utils/types.py
+++ b/tests/models/decoder_only/vision_language/vlm_utils/types.py
@@ -158,6 +158,7 @@ class VLMTestInfo(NamedTuple):
             "max_model_len": self.max_model_len,
             "max_num_seqs": self.max_num_seqs,
             "task": self.task,
+            "tensor_parallel_size": self.tensor_parallel_size,
             "hf_output_post_proc": self.hf_output_post_proc,
             "vllm_output_post_proc": self.vllm_output_post_proc,
             "auto_cls": self.auto_cls,