[XPU] fix all_reduce all-zero accuracy issue under torch.compile (#39844)
Signed-off-by: Chaojun Zhang <chaojun.zhang@intel.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
This commit is contained in:
co-authored by
mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
parent
48a65ccb02
commit
993859ceb0
@@ -47,9 +47,10 @@ class XpuCommunicator(DeviceCommunicatorBase):
|
||||
self.all2all_manager = AgRsAll2AllManager(self.cpu_group)
|
||||
logger.info("Using AgRs manager on XPU device.")
|
||||
|
||||
def all_reduce(self, input_) -> torch.Tensor:
|
||||
dist.all_reduce(input_, group=self.device_group)
|
||||
return input_
|
||||
def all_reduce(self, input_: torch.Tensor) -> torch.Tensor:
|
||||
output = input_.clone() if torch.compiler.is_compiling() else input_
|
||||
dist.all_reduce(output, group=self.device_group)
|
||||
return output
|
||||
|
||||
def reduce_scatter(self, input_: torch.Tensor, dim: int = -1):
|
||||
world_size = self.world_size
|
||||
|
||||
Reference in New Issue
Block a user