fastvideo.v1.distributed.device_communicators.cuda_communicator

fastvideo.v1.distributed.device_communicators.cuda_communicator#

Module Contents#

Classes#

CudaCommunicator

API#

class fastvideo.v1.distributed.device_communicators.cuda_communicator.CudaCommunicator(cpu_group: torch.distributed.ProcessGroup, device: Optional[torch.device] = None, device_group: Optional[torch.distributed.ProcessGroup] = None, unique_name: str = '')[source]#

Bases: fastvideo.v1.distributed.device_communicators.base_device_communicator.DeviceCommunicatorBase

all_reduce(input_, op: Optional[torch.distributed.ReduceOp] = None)[source]#
destroy() → None[source]#
recv(size: torch.Size, dtype: torch.dtype, src: Optional[int] = None) → torch.Tensor[source]#

Receives a tensor from the source rank.

send(tensor: torch.Tensor, dst: Optional[int] = None) → None[source]#

Sends a tensor to the destination rank in a non-blocking way

Contents