35 lines
1.1 KiB
Python
35 lines
1.1 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
from typing import Optional
|
|
|
|
import torch
|
|
from torch.distributed import ProcessGroup
|
|
|
|
from .base_device_communicator import DeviceCommunicatorBase
|
|
|
|
|
|
class CpuCommunicator(DeviceCommunicatorBase):
|
|
|
|
def __init__(self,
|
|
cpu_group: ProcessGroup,
|
|
device: Optional[torch.device] = None,
|
|
device_group: Optional[ProcessGroup] = None,
|
|
unique_name: str = ""):
|
|
super().__init__(cpu_group, device, device_group, unique_name)
|
|
self.ipex_available = False
|
|
self.dist_module = torch.distributed
|
|
try:
|
|
import intel_extension_for_pytorch as ipex
|
|
self.ipex_available = True
|
|
self.dist_module = ipex.distributed
|
|
except ImportError:
|
|
"""
|
|
Intel IPEX not found. Falling back to PyTorch native
|
|
all_reduce for CPU (e.g. MacOS)
|
|
"""
|
|
pass
|
|
|
|
def all_reduce(self, input_):
|
|
self.dist_module.all_reduce(input_, group=self.device_group)
|
|
return input_
|