Skip to content

cpu_communicator

Classes

fastvideo.distributed.device_communicators.cpu_communicator.CpuCommunicator 

CpuCommunicator(cpu_group: ProcessGroup, device: device | None = None, device_group: ProcessGroup | None = None, unique_name: str = '')

Bases: DeviceCommunicatorBase

Source code in fastvideo/distributed/device_communicators/cpu_communicator.py 
16
17
18
19
20
21
22
23
24
25
26
27
28
def __init__(self,
             cpu_group: ProcessGroup,
             device: torch.device | None = None,
             device_group: ProcessGroup | None = None,
             unique_name: str = ""):
    super().__init__(cpu_group, device, device_group, unique_name)
    self.dist_module = torch.distributed

    from fastvideo.platforms import current_platform

    if (current_platform.get_cpu_architecture() == CpuArchEnum.X86) and hasattr(
            torch.ops._C, "init_shm_manager") and unique_name.startswith("tp"):
        self.dist_module = _CPUSHMDistributed(self)

Methods:

fastvideo.distributed.device_communicators.cpu_communicator.CpuCommunicator.gather 
gather(input_: Tensor, dst: int = 0, dim: int = -1) -> Tensor | None

NOTE: We assume that the input tensor is on the same device across all the ranks. NOTE: dst is the local rank of the destination rank.

Source code in fastvideo/distributed/device_communicators/cpu_communicator.py 
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
def gather(self, input_: torch.Tensor, dst: int = 0, dim: int = -1) -> torch.Tensor | None:
    """
    NOTE: We assume that the input tensor is on the same device across
    all the ranks.
    NOTE: `dst` is the local rank of the destination rank.
    """
    world_size = self.world_size
    assert -input_.dim() <= dim < input_.dim(), (f"Invalid dim ({dim}) for input tensor with shape {input_.size()}")
    if dim < 0:
        # Convert negative dim to positive.
        dim += input_.dim()

    # Allocate output tensor.
    gather_list = [torch.empty_like(input_) for _ in range(world_size)] if self.rank_in_group == dst else None

    # Gather.
    self.dist_module.gather(input_, gather_list, dst=self.ranks[dst], group=self.device_group)

    output_tensor = torch.cat(gather_list, dim=dim) if self.rank_in_group == dst else None
    return output_tensor