[BugFix] Fix Embedding Models with TP>1 (#5075)

2024-05-28 08:32:42 -07:00 · 2024-05-28 08:32:42 -07:00 · 9ba415588a
commit 9ba415588a
parent d4f3985907
1 changed files with 4 additions and 0 deletions
--- a/vllm/worker/embedding_model_runner.py
+++ b/vllm/worker/embedding_model_runner.py
@ -79,6 +79,10 @@ class EmbeddingModelRunner(ModelRunner):
            execute_model_kwargs.update({"image_input": multi_modal_input})
        hidden_states = model_executable(**execute_model_kwargs)

+        # Only perform pooling in the driver worker.
+        if not self.is_driver_worker:
+            return None
+
        return self.model.pooler(hidden_states=hidden_states,
                                 pooling_metadata=pooling_metadata)