Python PyTorch SequenceEmbeddingAllToAll用法及代码示例

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

用法:

class torchrec.distributed.dist_data.SequenceEmbeddingAllToAll(pg: torch._C._distributed_c10d.ProcessGroup, features_per_rank: List[int], device: Optional[torch.device] = None)

参数：

• pg(dist.ProcessGroup) -AlltoAll 通信发生的进程组。

• features_per_rank(List[int]) -每个等级的特征数量列表。

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

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

根据拆分将序列嵌入重新分配到ProcessGroup。

例子：

init_distributed(rank=rank, size=2, backend="nccl")
pg = dist.new_group(backend="nccl")
features_per_rank = [4, 4]
m = SequenceEmbeddingAllToAll(pg, features_per_rank)
local_embs = torch.rand((6, 2))
sharding_ctx: SequenceShardingContext
output = m(
    local_embs=local_embs,
    lengths=sharding_ctx.lengths_after_input_dist,
    input_splits=sharding_ctx.input_splits,
    output_splits=sharding_ctx.output_splits,
    unbucketize_permute_tensor=None,
)
tensor = output.wait()