gahow/aituner
0
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Correct CollectiveSpec P0 phase accounting

Browse Source
This commit is contained in:
Gahow Wang
2026-07-13 18:08:52 +08:00
parent 7f4ae1708b
commit 08de0695e0
2 changed files with 295 additions and 45 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
scripts/collectivespec/p0_overlay/collectivespec_p0/worker.py
View File

@@ -25,6 +25,9 @@ class P0Worker(Worker):
original = model_runner._determine_batch_execution_and_padding
def traced_determine(*args: Any, **kwargs: Any): # type: ignore[no-untyped-def]
local_unpadded_rows = kwargs.get("num_tokens")
if local_unpadded_rows is None and len(args) > 1:
local_unpadded_rows = args[1]
output = original(*args, **kwargs)
self._p0_batch_epoch += 1
scheduler_output = kwargs.get("scheduler_output")
@@ -38,6 +41,7 @@ class P0Worker(Worker):
**self._p0_rank_fields(),
**plan_summary(scheduler_output),
cudagraph_mode=getattr(cudagraph_mode, "value", str(cudagraph_mode)),
local_unpadded_rows=local_unpadded_rows,
physical_batch_rows=getattr(batch_desc, "num_tokens", None),
should_ubatch=should_ubatch,
rows_across_dp=rows_across_dp,
@@ -54,7 +58,9 @@ class P0Worker(Worker):
"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),
# Worker.rank is the tensor-parallel rank within a DP replica in
# this vLLM executor. ParallelConfig does not expose it here.
"tensor_parallel_rank": getattr(self, "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),