[CI/Build] Fix VLM broadcast tests tensor_parallel_size passing (vl…

…lm-project#10161) Signed-off-by: Isotr0py <[email protected]> Signed-off-by: Maxime Fournioux <[email protected]>
mfournioux · Nov 20, 2024 · b22700d · b22700d
1 parent d17c4bd
commit b22700d
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 0 deletions.
diff --git a/tests/models/decoder_only/vision_language/test_models.py b/tests/models/decoder_only/vision_language/test_models.py
@@ -34,6 +34,7 @@
     "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
@@ -158,6 +158,7 @@ def get_non_parametrized_runner_kwargs(self):
             "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,