Python PyTorch PooledEmbeddingsAllToAll用法及代码示例

本文简要介绍python语言中 torchrec.distributed.dist_data.PooledEmbeddingsAllToAll 的用法。

用法:

class torchrec.distributed.dist_data.PooledEmbeddingsAllToAll(pg: torch._C._distributed_c10d.ProcessGroup, dim_sum_per_rank: List[int], device: Optional[torch.device] = None, callbacks: Optional[List[Callable[[torch.Tensor], torch.Tensor]]] = None)

参数：

• pg(dist.ProcessGroup) -ProcessGroup 用于AlltoAll 通信。

• dim_sum_per_rank(List[int]) -每个等级中嵌入的特征数(维度总和)。

• device(可选的[torch.device]) -将分配缓冲区的设备。

• callbacks(可选的[List[可调用[[torch.Tensor],torch.Tensor]]]) -

基础：torch.nn.modules.module.Module

根据 dim_sum_per_rank 对带有 ProcessGroup 的张量键进行分批和收集。

实现利用alltoall_pooled 操作。

例子：

dim_sum_per_rank = [2, 1]
a2a = PooledEmbeddingsAllToAll(pg, dim_sum_per_rank, device)

t0 = torch.rand((6, 2))
t1 = torch.rand((6, 1))
rank0_output = a2a(t0).wait()
rank1_output = a2a(t1).wait()
print(rank0_output.size())
    # torch.Size([3, 3])
print(rank1_output.size())
    # torch.Size([3, 3])