Revert "Fix different device type adjustment in PP" (#8141)

This commit is contained in:
Sai Enduri
2025-07-18 01:12:11 -07:00
committed by GitHub
parent 9d33fcfb8e
commit d0510f08fe
4 changed files with 27 additions and 25 deletions

View File

@@ -975,7 +975,6 @@ class Scheduler(
self.world_group.device_group,
self.pp_rank * self.tp_size + dp_offset,
(self.pp_rank + 1) * self.tp_size + dp_offset,
device=self.device,
)
# send out proxy tensors to the next stage
@@ -1024,7 +1023,6 @@ class Scheduler(
self.world_group.device_group,
(self.pp_rank - 1) * self.tp_size + dp_offset,
self.pp_rank * self.tp_size + dp_offset,
device=self.device,
)
else:
recv_reqs = None
@@ -1055,7 +1053,6 @@ class Scheduler(
self.attn_tp_group.rank,
self.attn_tp_cpu_group,
src=self.attn_tp_group.ranks[0],
device=self.device,
)
if self.tp_size != 1:
control_reqs = broadcast_pyobj(
@@ -1063,7 +1060,6 @@ class Scheduler(
self.tp_group.rank,
self.tp_cpu_group,
src=self.tp_group.ranks[0],
device=self.device,
)
recv_reqs = work_reqs + control_reqs
elif self.tp_size != 1:
@@ -1072,7 +1068,6 @@ class Scheduler(
self.tp_group.rank,
self.tp_cpu_group,
src=self.tp_group.ranks[0],
device=self.device,
)
return recv_reqs