Fix different device type adjustment in PP (#7760)

2025-07-15 19:37:14 -07:00
parent 7498522f7d
commit 3bc43c683e
4 changed files with 25 additions and 27 deletions
@@ -144,6 +144,7 @@ class TpModelWorker:
            self.tp_size * self.pp_rank + tp_rank,
            self.world_group.cpu_group,
            src=self.world_group.ranks[0],
+            device=self.device,
        )[0]
        set_random_seed(self.random_seed)