77 lines
3.1 KiB
Python
77 lines
3.1 KiB
Python
"""P0 worker: record actual rank-side target execution phases."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import Any
|
|
|
|
from vllm.v1.worker.gpu_worker import Worker
|
|
|
|
from .common import log_event, plan_summary
|
|
|
|
|
|
class P0Worker(Worker):
|
|
"""Log rank-local execution without modifying model or collective calls."""
|
|
|
|
def __init__(self, *args: Any, **kwargs: Any) -> None:
|
|
super().__init__(*args, **kwargs)
|
|
self._p0_execute_epoch = 0
|
|
self._p0_batch_epoch = 0
|
|
|
|
def init_device(self): # type: ignore[no-untyped-def]
|
|
result = super().init_device()
|
|
model_runner = self.model_runner
|
|
if getattr(model_runner, "_collectivespec_p0_wrapped", False):
|
|
return result
|
|
original = model_runner._determine_batch_execution_and_padding
|
|
|
|
def traced_determine(*args: Any, **kwargs: Any): # type: ignore[no-untyped-def]
|
|
output = original(*args, **kwargs)
|
|
self._p0_batch_epoch += 1
|
|
scheduler_output = kwargs.get("scheduler_output")
|
|
if scheduler_output is None and args:
|
|
scheduler_output = args[0]
|
|
cudagraph_mode, batch_desc, should_ubatch, rows_across_dp, _ = output
|
|
log_event(
|
|
"worker",
|
|
"batch_execution_plan",
|
|
batch_phase_epoch=self._p0_batch_epoch,
|
|
**self._p0_rank_fields(),
|
|
**plan_summary(scheduler_output),
|
|
cudagraph_mode=getattr(cudagraph_mode, "value", str(cudagraph_mode)),
|
|
physical_batch_rows=getattr(batch_desc, "num_tokens", None),
|
|
should_ubatch=should_ubatch,
|
|
rows_across_dp=rows_across_dp,
|
|
)
|
|
return output
|
|
|
|
model_runner._determine_batch_execution_and_padding = traced_determine
|
|
model_runner._collectivespec_p0_wrapped = True
|
|
return result
|
|
|
|
def _p0_rank_fields(self) -> dict[str, Any]:
|
|
config = self.vllm_config.parallel_config
|
|
return {
|
|
"global_rank": getattr(self, "rank", None),
|
|
"local_rank": getattr(self, "local_rank", None),
|
|
"data_parallel_rank": getattr(config, "data_parallel_rank", None),
|
|
"tensor_parallel_rank": getattr(config, "tensor_parallel_rank", None),
|
|
"expert_parallel_size": getattr(config, "expert_parallel_size", None),
|
|
"tensor_parallel_size": getattr(config, "tensor_parallel_size", None),
|
|
"data_parallel_size": getattr(config, "data_parallel_size", None),
|
|
"speculative_kmax": getattr(
|
|
getattr(self, "speculative_config", None), "num_speculative_tokens", None
|
|
),
|
|
}
|
|
|
|
def execute_model(self, scheduler_output: Any): # type: ignore[no-untyped-def]
|
|
self._p0_execute_epoch += 1
|
|
fields = {
|
|
"worker_phase_epoch": self._p0_execute_epoch,
|
|
**self._p0_rank_fields(),
|
|
**plan_summary(scheduler_output),
|
|
}
|
|
log_event("worker", "target_execute_begin", **fields)
|
|
output = super().execute_model(scheduler_output)
|
|
log_event("worker", "target_execute_end", **fields)
|
|
return output
|