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: torch.device | None = None, device_group: torch.distributed.ProcessGroup | None = None, unique_name: str = '')[source]#

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

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

Receives a tensor from the source rank.

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

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

Contents