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

9ba41558 · Robert Shaw · GitHub · d4f39859 · 9ba41558
Unverified Commit 9ba41558 authored May 28, 2024 by Robert Shaw Committed by GitHub May 28, 2024
Show whitespace changes
Inline Side-by-side

Showing with 4 additions and 0 deletions

vllm/worker/embedding_model_runner.py vllm/worker/embedding_model_runner.py +4 -0

No files found.
--- 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)