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