From 86465feb1270c99981837dc24e304f7069f0a985 Mon Sep 17 00:00:00 2001 From: Mahesh Keralapura Date: Thu, 1 Aug 2024 13:13:38 -0700 Subject: [PATCH 1/6] Add span metrics for model_forward, scheduler and sampler time --- tests/tracing/test_tracing.py | 2 ++ vllm/core/scheduler.py | 10 +++++++++- vllm/engine/llm_engine.py | 6 ++++++ vllm/sequence.py | 11 +++++++++++ vllm/tracing.py | 4 +++- vllm/worker/model_runner.py | 13 ++++++++++++- vllm/worker/model_runner_base.py | 7 +++++++ 7 files changed, 50 insertions(+), 3 deletions(-) diff --git a/tests/tracing/test_tracing.py b/tests/tracing/test_tracing.py index 2f8f62cf2..face29f8c 100644 --- a/tests/tracing/test_tracing.py +++ b/tests/tracing/test_tracing.py @@ -114,3 +114,5 @@ def test_traces(trace_service): SpanAttributes.LLM_LATENCY_TIME_TO_FIRST_TOKEN) == ttft e2e_time = metrics.finished_time - metrics.arrival_time assert attributes.get(SpanAttributes.LLM_LATENCY_E2E) == e2e_time + assert attributes.get( + SpanAttributes.LLM_LATENCY_TIME_IN_SCHEDULER) == metrics.scheduler_time_ms diff --git a/vllm/core/scheduler.py b/vllm/core/scheduler.py index 6e59c5e0f..fe8640129 100644 --- a/vllm/core/scheduler.py +++ b/vllm/core/scheduler.py @@ -975,6 +975,7 @@ def schedule(self) -> Tuple[List[SequenceGroupMetadata], SchedulerOutputs]: # Schedule sequence groups. # This function call changes the internal states of the scheduler # such as self.running, self.swapped, and self.waiting. + scheduler_start_time = time.time() scheduler_outputs = self._schedule() now = time.time() @@ -1047,7 +1048,14 @@ def schedule(self) -> Tuple[List[SequenceGroupMetadata], SchedulerOutputs]: for scheduled_seq_group in scheduler_outputs.scheduled_seq_groups: self.block_manager.mark_blocks_as_computed( scheduled_seq_group.seq_group) - + + scheduler_time = time.time() - scheduler_start_time + # Add this to scheduler time to all the sequences that are either currently running or + # swapped out. This is not added to the ones waiting on the queue and never scheduled. + # This will help estimate if the scheduler is a significant component in the e2e latency. + for seq_group in self.running + self.swapped: + seq_group.metrics.scheduler_time += scheduler_time + return seq_group_metadata_list, scheduler_outputs def fork_seq(self, parent_seq: Sequence, child_seq: Sequence) -> None: diff --git a/vllm/engine/llm_engine.py b/vllm/engine/llm_engine.py index eabe3b23a..1eb155528 100644 --- a/vllm/engine/llm_engine.py +++ b/vllm/engine/llm_engine.py @@ -814,6 +814,9 @@ def _process_model_outputs( seq_group = scheduled_seq_group.seq_group seq_group.update_num_computed_tokens( scheduled_seq_group.token_chunk_size) + for o in outputs + seq_group.metrics.model_forward_time += o.model_forward_time + seq_group.metrics.sampler_time += o.sampler_time if self.model_config.embedding_mode: self._process_sequence_group_outputs(seq_group, outputs) continue @@ -1205,3 +1208,6 @@ def create_trace_span(self, seq_group: SequenceGroup) -> None: seq_span.set_attribute( SpanAttributes.LLM_LATENCY_TIME_TO_FIRST_TOKEN, ttft) seq_span.set_attribute(SpanAttributes.LLM_LATENCY_E2E, e2e_time) + seq_span.set_attribute(SpanAttributes.LLM_LATENCY_TIME_IN_SCHEDULER, metrics.scheduler_time) + seq_span.set_attribute(SpanAttributes.LLM_LATENCY_TIME_IN_MODEL_FORWARD, metrics.model_forward_time) + seq_span.set_attribute(SpanAttributes.LLM_LATENCY_TIME_IN_SAMPLER, metrics.sampler_time) diff --git a/vllm/sequence.py b/vllm/sequence.py index 0cd4c7e71..d057e8ffe 100644 --- a/vllm/sequence.py +++ b/vllm/sequence.py @@ -91,6 +91,8 @@ class RequestMetrics: first_token_time: The time when the first token was generated. time_in_queue: The time the request spent in the queue. finished_time: The time when the request was finished. + scheduler_time: The time spent in the scheduler when this request was being considered by the scheduler. + model_forward_time: The time spent in the model forward pass when this request was in the batch. """ arrival_time: float last_token_time: float @@ -98,6 +100,9 @@ class RequestMetrics: first_token_time: Optional[float] time_in_queue: Optional[float] finished_time: Optional[float] = None + scheduler_time: Optional[float] = None + model_forward_time: Optional[float] = None + sampler_time: Optional[float] = None class SequenceData: @@ -863,6 +868,12 @@ class SamplerOutput: # Optional last hidden states from the model. hidden_states: Optional[torch.Tensor] = None + # Total time spent in the forward pass for this across all workers + model_forward_time: Optional[float] = None + + # Total time spent in the sampler. + sampler_time: Optional[float] = None + def __getitem__(self, idx: int): return self.outputs[idx] diff --git a/vllm/tracing.py b/vllm/tracing.py index ba6732cab..492ef3c07 100644 --- a/vllm/tracing.py +++ b/vllm/tracing.py @@ -92,7 +92,9 @@ class SpanAttributes(BaseSpanAttributes): LLM_LATENCY_TIME_IN_QUEUE = "gen_ai.latency.time_in_queue" LLM_LATENCY_TIME_TO_FIRST_TOKEN = "gen_ai.latency.time_to_first_token" LLM_LATENCY_E2E = "gen_ai.latency.e2e" - + LLM_LATENCY_TIME_IN_SCHEDULER = "gen_ai.latency.time_in_scheduler" + LLM_LATENCY_TIME_IN_MODEL_FORWARD = "gen_ai.latency.time_in_model_forward" + LLM_LATENCY_TIME_IN_SAMPLER = "gen_ai.latency.time_in_sampler" def contains_trace_headers(headers: Mapping[str, str]) -> bool: return any(h in headers for h in TRACE_HEADERS) diff --git a/vllm/worker/model_runner.py b/vllm/worker/model_runner.py index e63be184a..b5b9c5697 100644 --- a/vllm/worker/model_runner.py +++ b/vllm/worker/model_runner.py @@ -1311,6 +1311,9 @@ def execute_model( "finished_requests_ids": model_input.finished_requests_ids, "request_ids_to_seq_ids": model_input.request_ids_to_seq_ids, } if self.has_seqlen_agnostic else {} + model_forward_start = torch.cuda.Event(enable_timing=True) + model_forward_end = torch.cuda.Event(enable_timing=True) + model_forward_start.record() hidden_or_intermediate_states = model_executable( input_ids=model_input.input_tokens, positions=model_input.input_positions, @@ -1319,7 +1322,9 @@ def execute_model( intermediate_tensors=intermediate_tensors, **multi_modal_kwargs, **seqlen_agnostic_kwargs) - + model_forward_end.record() + model_forward_time = model_forward_start.elapsed_time(model_forward_end) + # Compute the logits in the last pipeline stage. if not get_pp_group().is_last_rank: return hidden_or_intermediate_states @@ -1330,11 +1335,17 @@ def execute_model( if not self.is_driver_worker: return [] + sampler_start = time.time() # Sample the next token. output: SamplerOutput = self.model.sample( logits=logits, sampling_metadata=model_input.sampling_metadata, ) + output.sampler_time = time.time() - sampler_start + # If there are multiple workers, we are still tracking the latency from the start time + # of the driver worker to the end time of the driver worker. The model forward time wil + # then end up covering the communication time as well. + output.model_forward_time = model_forward_time if self.return_hidden_states: # we only need to pass hidden states of most recent token diff --git a/vllm/worker/model_runner_base.py b/vllm/worker/model_runner_base.py index 5fb97025a..bdae08085 100644 --- a/vllm/worker/model_runner_base.py +++ b/vllm/worker/model_runner_base.py @@ -128,6 +128,13 @@ def build(self, *args, **kwargs) -> T: """Build metadata with on-device tensors.""" raise NotImplementedError +class ModelRunnerOutput: + """ + Model runner output that is used to collect the outputs and stats together. + """ + sampler_output: List[SamplerOutput] = [] + intermediate_tensors: IntermediateTensors = None + forward_time_ms: float = 0 class ModelRunnerBase(ABC, Generic[T]): """ From 2d90947aec21f8654f2f5faa18684dec9965c14e Mon Sep 17 00:00:00 2001 From: Mahesh Keralapura Date: Thu, 1 Aug 2024 17:13:57 -0700 Subject: [PATCH 2/6] Based on Aurick's suggestions - Collect model execute time instead of sampling time + get elapsed event time after sampler --- vllm/engine/llm_engine.py | 6 +++--- vllm/sequence.py | 7 ++++--- vllm/tracing.py | 2 +- vllm/worker/model_runner.py | 4 +--- vllm/worker/worker_base.py | 6 ++++-- 5 files changed, 13 insertions(+), 12 deletions(-) diff --git a/vllm/engine/llm_engine.py b/vllm/engine/llm_engine.py index 1eb155528..3b5be5853 100644 --- a/vllm/engine/llm_engine.py +++ b/vllm/engine/llm_engine.py @@ -814,9 +814,9 @@ def _process_model_outputs( seq_group = scheduled_seq_group.seq_group seq_group.update_num_computed_tokens( scheduled_seq_group.token_chunk_size) - for o in outputs + for o in outputs: seq_group.metrics.model_forward_time += o.model_forward_time - seq_group.metrics.sampler_time += o.sampler_time + seq_group.metrics.model_execute_time += o.model_execute_time if self.model_config.embedding_mode: self._process_sequence_group_outputs(seq_group, outputs) continue @@ -1210,4 +1210,4 @@ def create_trace_span(self, seq_group: SequenceGroup) -> None: seq_span.set_attribute(SpanAttributes.LLM_LATENCY_E2E, e2e_time) seq_span.set_attribute(SpanAttributes.LLM_LATENCY_TIME_IN_SCHEDULER, metrics.scheduler_time) seq_span.set_attribute(SpanAttributes.LLM_LATENCY_TIME_IN_MODEL_FORWARD, metrics.model_forward_time) - seq_span.set_attribute(SpanAttributes.LLM_LATENCY_TIME_IN_SAMPLER, metrics.sampler_time) + seq_span.set_attribute(SpanAttributes.LLM_LATENCY_TIME_IN_MODEL_EXECUTE, metrics.model_execute_time) diff --git a/vllm/sequence.py b/vllm/sequence.py index d057e8ffe..f516690dd 100644 --- a/vllm/sequence.py +++ b/vllm/sequence.py @@ -102,7 +102,7 @@ class RequestMetrics: finished_time: Optional[float] = None scheduler_time: Optional[float] = None model_forward_time: Optional[float] = None - sampler_time: Optional[float] = None + model_execute_time: Optional[float] = None class SequenceData: @@ -871,8 +871,9 @@ class SamplerOutput: # Total time spent in the forward pass for this across all workers model_forward_time: Optional[float] = None - # Total time spent in the sampler. - sampler_time: Optional[float] = None + # Total time spent in the model execute function. This will include model forward, + # block/sync across workers, cpu-gpu sync time and sampling time. + model_execute_time: Optional[float] = None def __getitem__(self, idx: int): return self.outputs[idx] diff --git a/vllm/tracing.py b/vllm/tracing.py index 492ef3c07..556ec8623 100644 --- a/vllm/tracing.py +++ b/vllm/tracing.py @@ -94,7 +94,7 @@ class SpanAttributes(BaseSpanAttributes): LLM_LATENCY_E2E = "gen_ai.latency.e2e" LLM_LATENCY_TIME_IN_SCHEDULER = "gen_ai.latency.time_in_scheduler" LLM_LATENCY_TIME_IN_MODEL_FORWARD = "gen_ai.latency.time_in_model_forward" - LLM_LATENCY_TIME_IN_SAMPLER = "gen_ai.latency.time_in_sampler" + LLM_LATENCY_TIME_IN_MODEL_EXECUTE = "gen_ai.latency.time_in_model_execute" def contains_trace_headers(headers: Mapping[str, str]) -> bool: return any(h in headers for h in TRACE_HEADERS) diff --git a/vllm/worker/model_runner.py b/vllm/worker/model_runner.py index b5b9c5697..e72d27134 100644 --- a/vllm/worker/model_runner.py +++ b/vllm/worker/model_runner.py @@ -1323,7 +1323,6 @@ def execute_model( **multi_modal_kwargs, **seqlen_agnostic_kwargs) model_forward_end.record() - model_forward_time = model_forward_start.elapsed_time(model_forward_end) # Compute the logits in the last pipeline stage. if not get_pp_group().is_last_rank: @@ -1335,13 +1334,12 @@ def execute_model( if not self.is_driver_worker: return [] - sampler_start = time.time() # Sample the next token. output: SamplerOutput = self.model.sample( logits=logits, sampling_metadata=model_input.sampling_metadata, ) - output.sampler_time = time.time() - sampler_start + model_forward_time = model_forward_start.elapsed_time(model_forward_end) # If there are multiple workers, we are still tracking the latency from the start time # of the driver worker to the end time of the driver worker. The model forward time wil # then end up covering the communication time as well. diff --git a/vllm/worker/worker_base.py b/vllm/worker/worker_base.py index 03e3857e2..287a69fec 100644 --- a/vllm/worker/worker_base.py +++ b/vllm/worker/worker_base.py @@ -2,6 +2,7 @@ import importlib import os from abc import ABC, abstractmethod +import time from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Type, Union import torch @@ -217,6 +218,7 @@ def execute_model( self, execute_model_req: Optional[ExecuteModelRequest] = None ) -> Optional[List[SamplerOutput]]: + start_time = time.time() """Executes at least one model step on the given sequences, unless no sequences are provided.""" if self.is_driver_worker: @@ -268,12 +270,12 @@ def execute_model( if not get_pp_group().is_first_rank: intermediate_tensors = IntermediateTensors( get_pp_group().recv_tensor_dict()) - + output = self.model_runner.execute_model( model_input, self.kv_cache[worker_input.virtual_engine] if self.kv_cache is not None else None, intermediate_tensors, num_steps) - + output[0].model_execute_time = time.time() - start_time if not get_pp_group().is_last_rank: # output is IntermediateTensors get_pp_group().send_tensor_dict(output.tensors) From f53c88c3b3e76e79d4b066c6c236bc7991ba26bc Mon Sep 17 00:00:00 2001 From: Mahesh Keralapura Date: Thu, 1 Aug 2024 17:56:44 -0700 Subject: [PATCH 3/6] Fix bug: Handle the case where the metrics have not yet been initialized --- vllm/core/scheduler.py | 5 ++++- vllm/engine/llm_engine.py | 13 ++++++++++--- vllm/worker/worker_base.py | 5 ++++- 3 files changed, 18 insertions(+), 5 deletions(-) diff --git a/vllm/core/scheduler.py b/vllm/core/scheduler.py index fe8640129..eacaf1078 100644 --- a/vllm/core/scheduler.py +++ b/vllm/core/scheduler.py @@ -1054,7 +1054,10 @@ def schedule(self) -> Tuple[List[SequenceGroupMetadata], SchedulerOutputs]: # swapped out. This is not added to the ones waiting on the queue and never scheduled. # This will help estimate if the scheduler is a significant component in the e2e latency. for seq_group in self.running + self.swapped: - seq_group.metrics.scheduler_time += scheduler_time + if seq_group.metrics.scheduler_time is not None: + seq_group.metrics.scheduler_time += scheduler_time + else: + seq_group.metrics.scheduler_time = scheduler_time return seq_group_metadata_list, scheduler_outputs diff --git a/vllm/engine/llm_engine.py b/vllm/engine/llm_engine.py index 3b5be5853..34f8da4b4 100644 --- a/vllm/engine/llm_engine.py +++ b/vllm/engine/llm_engine.py @@ -814,9 +814,16 @@ def _process_model_outputs( seq_group = scheduled_seq_group.seq_group seq_group.update_num_computed_tokens( scheduled_seq_group.token_chunk_size) - for o in outputs: - seq_group.metrics.model_forward_time += o.model_forward_time - seq_group.metrics.model_execute_time += o.model_execute_time + if output is not None and len(output) > 0 and isinstance(output[0], SamplerOutput): + for o in output: + if seq_group.metrics.model_forward_time is not None: + seq_group.metrics.model_forward_time += o.model_forward_time + else: + seq_group.metrics.model_forward_time = o.model_forward_time + if seq_group.metrics.model_execute_time is not None: + seq_group.metrics.model_execute_time += o.model_execute_time + else: + seq_group.metrics.model_execute_time = o.model_execute_time if self.model_config.embedding_mode: self._process_sequence_group_outputs(seq_group, outputs) continue diff --git a/vllm/worker/worker_base.py b/vllm/worker/worker_base.py index 287a69fec..b7381372c 100644 --- a/vllm/worker/worker_base.py +++ b/vllm/worker/worker_base.py @@ -275,7 +275,10 @@ def execute_model( model_input, self.kv_cache[worker_input.virtual_engine] if self.kv_cache is not None else None, intermediate_tensors, num_steps) - output[0].model_execute_time = time.time() - start_time + end_time = time.time() + if output is not None: + for o in output: + o.model_execute_time = end_time - start_time if not get_pp_group().is_last_rank: # output is IntermediateTensors get_pp_group().send_tensor_dict(output.tensors) From 36f876c8d06e63403a4c02a3d9b702558632dbe6 Mon Sep 17 00:00:00 2001 From: Mahesh Keralapura Date: Fri, 2 Aug 2024 08:10:33 -0700 Subject: [PATCH 4/6] Make the tracing of model forward time optional behind a command line arg --- vllm/config.py | 5 +++++ vllm/engine/arg_utils.py | 12 +++++++++++- vllm/engine/llm_engine.py | 10 +++++++--- vllm/executor/executor_base.py | 7 ++++--- vllm/executor/gpu_executor.py | 1 + vllm/worker/model_runner.py | 17 ++++++++++------- vllm/worker/worker.py | 7 +++++-- 7 files changed, 43 insertions(+), 16 deletions(-) diff --git a/vllm/config.py b/vllm/config.py index 6f0fdf8bc..fe0487ade 100644 --- a/vllm/config.py +++ b/vllm/config.py @@ -1576,6 +1576,11 @@ class ObservabilityConfig: """Configuration for observability.""" otlp_traces_endpoint: Optional[str] = None + # If set, collects the model forward time for the request. This introduces a possibly + # blocking operation to accurately collect the GPU time. It can have a performance + # impact on the request latency. + collect_model_forward_time: bool = False + def __post_init__(self): if not is_otel_installed() and self.otlp_traces_endpoint is not None: raise ValueError("OpenTelemetry packages must be installed before " diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index 05bfe7c24..de72f5a8f 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -117,6 +117,7 @@ class EngineArgs: disable_logprobs_during_spec_decoding: Optional[bool] = None otlp_traces_endpoint: Optional[str] = None + collect_model_forward_time: bool = False def __post_init__(self): if self.tokenizer is None: @@ -660,6 +661,14 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser: type=str, default=None, help='Target URL to which OpenTelemetry traces will be sent.') + parser.add_argument( + '--collect-model-forward-time', + type=bool, + default=EngineArgs.collect_model_forward_time, + help="If set to True and otlp-traces-endpoint is set, " + "collects model forward time in traces. This involves " + "use of a blocking operation and hence might have a " + "performance impact.") return parser @@ -845,7 +854,8 @@ def create_engine_config(self, ) -> EngineConfig: guided_decoding_backend=self.guided_decoding_backend) observability_config = ObservabilityConfig( - otlp_traces_endpoint=self.otlp_traces_endpoint) + otlp_traces_endpoint=self.otlp_traces_endpoint, + collect_model_forward_time=self.collect_model_forward_time) if (model_config.get_sliding_window() is not None and scheduler_config.chunked_prefill_enabled diff --git a/vllm/engine/llm_engine.py b/vllm/engine/llm_engine.py index 34f8da4b4..888d0e00c 100644 --- a/vllm/engine/llm_engine.py +++ b/vllm/engine/llm_engine.py @@ -259,6 +259,7 @@ def __init__( speculative_config=speculative_config, load_config=load_config, prompt_adapter_config=prompt_adapter_config, + observability_config=self.observability_config, ) if not self.model_config.embedding_mode: @@ -1215,6 +1216,9 @@ def create_trace_span(self, seq_group: SequenceGroup) -> None: seq_span.set_attribute( SpanAttributes.LLM_LATENCY_TIME_TO_FIRST_TOKEN, ttft) seq_span.set_attribute(SpanAttributes.LLM_LATENCY_E2E, e2e_time) - seq_span.set_attribute(SpanAttributes.LLM_LATENCY_TIME_IN_SCHEDULER, metrics.scheduler_time) - seq_span.set_attribute(SpanAttributes.LLM_LATENCY_TIME_IN_MODEL_FORWARD, metrics.model_forward_time) - seq_span.set_attribute(SpanAttributes.LLM_LATENCY_TIME_IN_MODEL_EXECUTE, metrics.model_execute_time) + if metrics.scheduler_time is not None: + seq_span.set_attribute(SpanAttributes.LLM_LATENCY_TIME_IN_SCHEDULER, metrics.scheduler_time) + if metrics.model_forward_time is not None: + seq_span.set_attribute(SpanAttributes.LLM_LATENCY_TIME_IN_MODEL_FORWARD, metrics.model_forward_time / 1000.0) + if metrics.model_execute_time is not None: + seq_span.set_attribute(SpanAttributes.LLM_LATENCY_TIME_IN_MODEL_EXECUTE, metrics.model_execute_time) diff --git a/vllm/executor/executor_base.py b/vllm/executor/executor_base.py index a848bc709..bc4f54455 100644 --- a/vllm/executor/executor_base.py +++ b/vllm/executor/executor_base.py @@ -2,8 +2,8 @@ from typing import List, Optional, Set, Tuple from vllm.config import (CacheConfig, DeviceConfig, LoadConfig, LoRAConfig, - ModelConfig, MultiModalConfig, ParallelConfig, - PromptAdapterConfig, SchedulerConfig, + ModelConfig, MultiModalConfig, ObservabilityConfig, + ParallelConfig, PromptAdapterConfig, SchedulerConfig, SpeculativeConfig) from vllm.lora.request import LoRARequest from vllm.prompt_adapter.request import PromptAdapterRequest @@ -32,6 +32,7 @@ def __init__( multimodal_config: Optional[MultiModalConfig], speculative_config: Optional[SpeculativeConfig], prompt_adapter_config: Optional[PromptAdapterConfig], + observability_config: Optional[ObservabilityConfig], ) -> None: self.model_config = model_config self.cache_config = cache_config @@ -43,7 +44,7 @@ def __init__( self.multimodal_config = multimodal_config self.speculative_config = speculative_config self.prompt_adapter_config = prompt_adapter_config - + self.observability_config = observability_config self._init_executor() @abstractmethod diff --git a/vllm/executor/gpu_executor.py b/vllm/executor/gpu_executor.py index 3e77af0e2..57b9e2b33 100644 --- a/vllm/executor/gpu_executor.py +++ b/vllm/executor/gpu_executor.py @@ -60,6 +60,7 @@ def _get_worker_kwargs( prompt_adapter_config=self.prompt_adapter_config, is_driver_worker=(not self.parallel_config) or (rank % self.parallel_config.tensor_parallel_size == 0), + observability_config=self.observability_config, ) def _get_create_worker_kwargs( diff --git a/vllm/worker/model_runner.py b/vllm/worker/model_runner.py index e72d27134..f366978c1 100644 --- a/vllm/worker/model_runner.py +++ b/vllm/worker/model_runner.py @@ -25,8 +25,8 @@ from vllm.attention import AttentionMetadata, get_attn_backend from vllm.config import (CacheConfig, DeviceConfig, LoadConfig, LoRAConfig, - ModelConfig, MultiModalConfig, ParallelConfig, - PromptAdapterConfig, SchedulerConfig) + ModelConfig, MultiModalConfig, ObservabilityConfig, + ParallelConfig, PromptAdapterConfig, SchedulerConfig) from vllm.distributed import get_pp_group from vllm.distributed.parallel_state import graph_capture from vllm.inputs import INPUT_REGISTRY @@ -607,6 +607,7 @@ def __init__( prompt_adapter_config: Optional[PromptAdapterConfig] = None, multimodal_config: Optional[MultiModalConfig] = None, return_hidden_states: bool = False, + observability_config: Optional[ObservabilityConfig] = None, ): self.model_config = model_config self.parallel_config = parallel_config @@ -619,6 +620,7 @@ def __init__( self.prompt_adapter_config = prompt_adapter_config self.multimodal_config = multimodal_config self.return_hidden_states = return_hidden_states + self.observability_config = observability_config self.device = self.device_config.device self.pin_memory = is_pin_memory_available() @@ -1339,11 +1341,12 @@ def execute_model( logits=logits, sampling_metadata=model_input.sampling_metadata, ) - model_forward_time = model_forward_start.elapsed_time(model_forward_end) - # If there are multiple workers, we are still tracking the latency from the start time - # of the driver worker to the end time of the driver worker. The model forward time wil - # then end up covering the communication time as well. - output.model_forward_time = model_forward_time + if self.observability_config.collect_model_forward_time: + model_forward_time = model_forward_start.elapsed_time(model_forward_end) + # If there are multiple workers, we are still tracking the latency from the start time + # of the driver worker to the end time of the driver worker. The model forward time wil + # then end up covering the communication time as well. + output.model_forward_time = model_forward_time if self.return_hidden_states: # we only need to pass hidden states of most recent token diff --git a/vllm/worker/worker.py b/vllm/worker/worker.py index f3c379d1a..c88c28d57 100644 --- a/vllm/worker/worker.py +++ b/vllm/worker/worker.py @@ -7,8 +7,8 @@ import torch.distributed from vllm.config import (CacheConfig, DeviceConfig, LoadConfig, LoRAConfig, - ModelConfig, MultiModalConfig, ParallelConfig, - PromptAdapterConfig, SchedulerConfig, + ModelConfig, MultiModalConfig, ObservabilityConfig, + ParallelConfig, PromptAdapterConfig, SchedulerConfig, SpeculativeConfig) from vllm.distributed import (ensure_model_parallel_initialized, init_distributed_environment, @@ -50,6 +50,7 @@ def __init__( prompt_adapter_config: Optional[PromptAdapterConfig] = None, is_driver_worker: bool = False, model_runner_cls: Optional[Type[GPUModelRunnerBase]] = None, + observability_config: Optional[ObservabilityConfig] = None, ) -> None: self.model_config = model_config self.parallel_config = parallel_config @@ -72,6 +73,7 @@ def __init__( from vllm.utils import init_cached_hf_modules init_cached_hf_modules() self.multimodal_config = multimodal_config + self.observability_config = observability_config # Return hidden states from target model if the draft model is an # mlp_speculator @@ -99,6 +101,7 @@ def __init__( is_driver_worker=is_driver_worker, prompt_adapter_config=prompt_adapter_config, multimodal_config=multimodal_config, + observability_config=observability_config, **speculative_args, ) # Uninitialized cache engine. Will be initialized by From 3a2f2c1bf08dc4ba9187335637c17179d86c313d Mon Sep 17 00:00:00 2001 From: Mahesh Keralapura Date: Fri, 2 Aug 2024 10:57:47 -0700 Subject: [PATCH 5/6] use sync --- vllm/worker/model_runner.py | 1 + 1 file changed, 1 insertion(+) diff --git a/vllm/worker/model_runner.py b/vllm/worker/model_runner.py index f366978c1..fc56c75dd 100644 --- a/vllm/worker/model_runner.py +++ b/vllm/worker/model_runner.py @@ -1342,6 +1342,7 @@ def execute_model( sampling_metadata=model_input.sampling_metadata, ) if self.observability_config.collect_model_forward_time: + model_forward_end.synchronize() model_forward_time = model_forward_start.elapsed_time(model_forward_end) # If there are multiple workers, we are still tracking the latency from the start time # of the driver worker to the end time of the driver worker. The model forward time wil From c836bca15870e33c78b5cc4a10e71afde6df9e7a Mon Sep 17 00:00:00 2001 From: Mahesh Keralapura Date: Fri, 2 Aug 2024 13:02:57 -0700 Subject: [PATCH 6/6] Remove swapped from the set over which we count the scheduler time --- vllm/core/scheduler.py | 5 ++--- vllm/worker/worker_base.py | 2 +- 2 files changed, 3 insertions(+), 4 deletions(-) diff --git a/vllm/core/scheduler.py b/vllm/core/scheduler.py index eacaf1078..44924ef4d 100644 --- a/vllm/core/scheduler.py +++ b/vllm/core/scheduler.py @@ -1050,10 +1050,9 @@ def schedule(self) -> Tuple[List[SequenceGroupMetadata], SchedulerOutputs]: scheduled_seq_group.seq_group) scheduler_time = time.time() - scheduler_start_time - # Add this to scheduler time to all the sequences that are either currently running or - # swapped out. This is not added to the ones waiting on the queue and never scheduled. + # Add this to scheduler time to all the sequences that are currently running. # This will help estimate if the scheduler is a significant component in the e2e latency. - for seq_group in self.running + self.swapped: + for seq_group in self.running: if seq_group.metrics.scheduler_time is not None: seq_group.metrics.scheduler_time += scheduler_time else: diff --git a/vllm/worker/worker_base.py b/vllm/worker/worker_base.py index b7381372c..2a0693e6b 100644 --- a/vllm/worker/worker_base.py +++ b/vllm/worker/worker_base.py @@ -270,7 +270,7 @@ def execute_model( if not get_pp_group().is_first_rank: intermediate_tensors = IntermediateTensors( get_pp_group().recv_tensor_dict()) - + output = self.model_runner.execute_model( model_input, self.kv_cache[worker_input.virtual_engine] if self.kv_cache is not None else None, intermediate_tensors,