[Bugfix][CPU] Fix CPU embedding runner with tensor parallel (#10394)

Signed-off-by: Isotr0py <[email protected]>
vllm-project · Nov 17, 2024 · cf349c4 · cf349c4
1 parent 905d0f0
commit cf349c4
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/vllm/worker/cpu_embedding_model_runner.py b/vllm/worker/cpu_embedding_model_runner.py
@@ -66,6 +66,10 @@ def execute_model(
 
         hidden_states = model_executable(**execute_model_kwargs)
 
+        # Only perform pooling in the driver worker.
+        if not self.is_driver_worker:
+            return []
+
         return [
             self.model.pooler(hidden_states=hidden_states,
                               pooling_metadata=model_input.pooling_metadata)