[Cherry-pick][v0.18.0.post1][Bugfix] Fix VAE parallelism dist.gather performance bottleneck on NPU (#2970)

Signed-off-by: fan2956 <zhoufan53@huawei.com>
This commit is contained in:
fan2956 authored and GitHub committed 2026-04-21 20:36:08 +08:00
1 parent 89f733d941
commit caa26e1d04
1 file changed
+3 -3
@@ -54,9 +54,9 @@ class DistributedVaeExecutor:
self.parallel_size = parallel_size
def gather_tensors(self, tensor: torch.Tensor):
gather_list = [torch.empty_like(tensor) for _ in range(self.world_size)] if self.rank == 0 else None
dist.gather(tensor, gather_list=gather_list, dst=0, group=self.group)
return gather_list
gather_list = [torch.empty_like(tensor) for _ in range(self.world_size)]
dist.all_gather(gather_list, tensor, group=self.group)
return gather_list if self.rank == 0 else None
def broadcast_tensor(self, tensor: torch.Tensor):
dist.broadcast(tensor, src=0, group=self.group)