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])