[bug fix][pp] fix inconsistent latency between tp (#15379)

Signed-off-by: Xuchun Shang <xuchun.shang@gmail.com>
Co-authored-by: Shangming Cai <csmthu@gmail.com>
This commit is contained in:
Xuchun Shang
2025-12-18 16:59:56 +08:00
committed by GitHub
parent 56d12b4aea
commit fea2d5211d

View File

@@ -640,6 +640,16 @@ class SchedulerPPMixin:
f"seq_lens={seq_lens}, latencies_ms={latencies}"
)
if self.attn_tp_size > 1:
data_to_sync_tp = [seq_lens, latencies]
data_to_sync_tp = broadcast_pyobj(
data_to_sync_tp,
self.attn_tp_group.rank,
self.attn_tp_cpu_group,
src=self.attn_tp_group.ranks[0],
)
seq_lens, latencies = data_to_sync_tp
# Broadcast data to all ranks
if torch.distributed.is_available() and torch.distributed.is_initialized():
data_to_sync = [seq_lens, latencies]
@@ -869,7 +879,7 @@ class SchedulerPPMixin:
else:
data = None
if self.attn_tp_size != 1:
if self.attn_tp_size > 1:
data = broadcast_pyobj(
data,
self.attn_tp_group.rank,