fix tp

Signed-off-by: Lu Fang <fanglu@fb.com>

Author: luccafong

vllm/v1/engine/core.py

@@ -43,9 +43,6 @@
 from vllm.v1.serial_utils import MsgpackDecoder, MsgpackEncoder
 from vllm.v1.structured_output import StructuredOutputManager
 from vllm.version import __version__ as VLLM_VERSION
-from vllm.v1.worker.intermediates_logging import (
-    intermediate_logging
-)
 
 logger = init_logger(__name__)
 
@@ -221,8 +218,7 @@ def abort_requests(self, request_ids: list[str]):
     def execute_model(self, scheduler_output: SchedulerOutput):
         try:
             # Execute the model
-            with intermediate_logging(self.vllm_config.il_config):
-                return self.model_executor.execute_model(scheduler_output)
+            return self.model_executor.execute_model(scheduler_output)
         except Exception as err:
             # We do not want to catch BaseException here since we're only
             # interested in dumping info when the exception is due to an

vllm/v1/worker/gpu_worker.py

@@ -30,6 +30,7 @@
 from vllm.v1.utils import report_usage_stats
 from vllm.v1.worker.gpu_model_runner import GPUModelRunner
 from vllm.v1.worker.worker_base import WorkerBase
+from vllm.v1.worker.intermediates_logging import intermediate_logging
 
 logger = init_logger(__name__)
 
@@ -320,8 +321,9 @@ def execute_model(
                 get_pp_group().recv_tensor_dict(
                     all_gather_group=get_tp_group()))
 
-        output = self.model_runner.execute_model(scheduler_output,
-                                                 intermediate_tensors)
+        with intermediate_logging(self.vllm_config.il_config):
+            output = self.model_runner.execute_model(scheduler_output,
+                                                     intermediate_tensors)
 
         parallel_config = self.vllm_config.parallel_config
         if parallel_config.distributed_executor_backend != "external_launcher" \