Python PyTorch PooledEmbeddingsReduceScatter用法及代码示例

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

用法:

class torchrec.distributed.dist_data.PooledEmbeddingsReduceScatter(pg: torch._C._distributed_c10d.ProcessGroup)

参数：

pg(dist.ProcessGroup) -reduce-scatter 通信发生在其中的进程组。

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

包装 reduce-scatter 通信原语的模块类，用于按行和 twrw 分片中的池化嵌入通信。

对于池化嵌入，我们有一个局部 model-parallel 输出张量，其布局为 [num_buckets x batch_size，维度]。我们需要对各个批次的num_buckets 维度求和。我们沿着第一维将张量分成相等的块(不同桶的张量切片)，并将它们减少为输出张量，并将结果分散到相应的等级。

该类返回池化嵌入张量的异步 Awaitable 句柄。 reduce-scatter 仅适用于 NCCL 后端。

例子：

init_distributed(rank=rank, size=2, backend="nccl")
pg = dist.new_group(backend="nccl")
input = torch.randn(2 * 2, 2)
m = PooledEmbeddingsReduceScatter(pg)
output = m(input)
tensor = output.wait()