From a1271fa35dcf480e48e9001125110cfd4f5ce1bf Mon Sep 17 00:00:00 2001 From: Lucas Wilkinson Date: Mon, 16 Sep 2024 21:11:11 +0000 Subject: [PATCH] fix duplicate entries in the tp > 1 case --- vllm/profiler/layerwise_profile.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/vllm/profiler/layerwise_profile.py b/vllm/profiler/layerwise_profile.py index e074cc3f0d2ad..9d9f427e807f6 100644 --- a/vllm/profiler/layerwise_profile.py +++ b/vllm/profiler/layerwise_profile.py @@ -161,6 +161,11 @@ def _build_module_tree(self): def _df_traversal(event: _ProfilerEvent, curr_node: Optional[_ModuleTreeNode] = None): + + # For the tensor parallel case for now only look at task 1 + if event.start_tid != 1: + return + if event_has_module(event): node = _ModuleTreeNode(event=event, parent=curr_node) if curr_node: